nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: cubed-dev/cubed
Pull request #415 (2024-03-06 18:00:47)

Mark stack operation as not fusable
Simple fix for the problem described in #414 until the full fix is implemented.


Patch:
diff --git a/cubed/array_api/manipulation_functions.py b/cubed/array_api/manipulation_functions.py
index 48e0fbf3..945fcc66 100644
--- a/cubed/array_api/manipulation_functions.py
+++ b/cubed/array_api/manipulation_functions.py
@@ -296,6 +296,10 @@ def stack(arrays, /, *, axis=0):
         in_name = array_names[out_coords[axis]]
         return ((in_name, *(out_coords[:axis] + out_coords[(axis + 1) :])),)
 
+    # We have to mark this as fusable=False since the number of input args to
+    # the _read_stack_chunk function is *not* the same as the number of
+    # predecessor nodes in the DAG, and the fusion functions in blockwise
+    # assume they are the same. See https://github.com/cubed-dev/cubed/issues/414
     return general_blockwise(
         _read_stack_chunk,
         block_function,
@@ -304,6 +308,7 @@ def stack(arrays, /, *, axis=0):
         dtype=dtype,
         chunks=chunks,
         axis=axis,
+        fusable=False,
     )
 
 


Test patch:
diff --git a/cubed/tests/test_optimization.py b/cubed/tests/test_optimization.py
index 88d9b17f..3f8e9fc7 100644
--- a/cubed/tests/test_optimization.py
+++ b/cubed/tests/test_optimization.py
@@ -7,6 +7,7 @@ from numpy.testing import assert_array_equal
 
 import cubed
 import cubed.array_api as xp
+import cubed.random
 from cubed.backend_array_api import namespace as nxp
 from cubed.core.ops import elemwise, merge_chunks_new, partial_reduce
 from cubed.core.optimization import (
@@ -907,3 +908,12 @@ def test_fuse_only_optimize_dag(spec):
 
     result = d.compute(optimize_function=opt_fn)
     assert_array_equal(result, -np.ones((2, 2)))
+
+
+def test_optimize_stack(spec):
+    # This test fails if stack's general_blockwise call doesn't have fusable=False
+    a = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
+    b = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
+    c = xp.stack((a, b), axis=0)
+    # try to fuse all ops into one
+    c.compute(optimize_function=fuse_multiple_levels(max_total_num_input_blocks=10))
Repository: cubed-dev/cubed
Pull request #433 (2024-03-19 15:38:36)

Change `total_nbytes` to `total_nbytes_written`
This is to avoid reporting the size of Zarr arrays being read from (using `from_zarr`) in the intermediate `total_nbytes`.


Patch:
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
index 86d63c15..ba8b64fa 100644
--- a/cubed/core/plan.py
+++ b/cubed/core/plan.py
@@ -245,14 +245,16 @@ class Plan:
         ]
         return max(projected_mem_values) if len(projected_mem_values) > 0 else 0
 
-    def total_nbytes(self, optimize_graph: bool = True, optimize_function=None) -> int:
-        """Return the total number of bytes for all materialized arrays in this plan."""
+    def total_nbytes_written(
+        self, optimize_graph: bool = True, optimize_function=None
+    ) -> int:
+        """Return the total number of bytes written for all materialized arrays in this plan."""
         dag = self._finalize_dag(optimize_graph, optimize_function)
         nbytes = 0
         for _, d in dag.nodes(data=True):
             if d.get("type") == "array":
                 target = d["target"]
-                if isinstance(target, (LazyZarrArray, zarr.Array)):
+                if isinstance(target, LazyZarrArray):
                     nbytes += target.nbytes
         return nbytes
 
@@ -282,7 +284,7 @@ class Plan:
                 # note that \l is used to left-justify each line (see https://www.graphviz.org/docs/attrs/nojustify/)
                 rf"num tasks: {self.num_tasks(optimize_graph, optimize_function)}\l"
                 rf"max projected memory: {memory_repr(self.max_projected_mem(optimize_graph, optimize_function))}\l"
-                rf"total nbytes: {memory_repr(self.total_nbytes(optimize_graph, optimize_function))}\l"
+                rf"total nbytes written: {memory_repr(self.total_nbytes_written(optimize_graph, optimize_function))}\l"
                 rf"optimized: {optimize_graph}\l"
             ),
             "labelloc": "bottom",


Test patch:
diff --git a/cubed/tests/test_optimization.py b/cubed/tests/test_optimization.py
index 56a12d33..abc4b0a0 100644
--- a/cubed/tests/test_optimization.py
+++ b/cubed/tests/test_optimization.py
@@ -37,12 +37,15 @@ def test_fusion(spec):
     num_created_arrays = 3  # b, c, d (a is not created on disk)
     assert d.plan.num_arrays(optimize_graph=False) == num_arrays
     assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 12
-    assert d.plan.total_nbytes(optimize_graph=False) == b.nbytes + c.nbytes + d.nbytes
+    assert (
+        d.plan.total_nbytes_written(optimize_graph=False)
+        == b.nbytes + c.nbytes + d.nbytes
+    )
     num_arrays = 2  # a, d
     num_created_arrays = 1  # d (a is not created on disk)
     assert d.plan.num_arrays(optimize_graph=True) == num_arrays
     assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 4
-    assert d.plan.total_nbytes(optimize_graph=True) == d.nbytes
+    assert d.plan.total_nbytes_written(optimize_graph=True) == d.nbytes
 
     task_counter = TaskCounter()
     result = d.compute(callbacks=[task_counter])
Repository: cubed-dev/cubed
Pull request #532 (2024-08-01 11:23:30)

Add better test for local paths in `delete_on_exit`
Follow up to #524 to cover relative paths


Patch:
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
index 919bd14c..e91070bf 100644
--- a/cubed/core/plan.py
+++ b/cubed/core/plan.py
@@ -16,7 +16,13 @@ from cubed.primitive.types import PrimitiveOperation
 from cubed.runtime.pipeline import visit_nodes
 from cubed.runtime.types import ComputeEndEvent, ComputeStartEvent, CubedPipeline
 from cubed.storage.zarr import LazyZarrArray
-from cubed.utils import chunk_memory, extract_stack_summaries, join_path, memory_repr
+from cubed.utils import (
+    chunk_memory,
+    extract_stack_summaries,
+    is_local_path,
+    join_path,
+    memory_repr,
+)
 
 # A unique ID with sensible ordering, used for making directory names
 CONTEXT_ID = f"cubed-{datetime.now().strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4()}"
@@ -26,7 +32,7 @@ CONTEXT_DIRS = set()
 
 
 def delete_on_exit(context_dir: str) -> None:
-    if context_dir not in CONTEXT_DIRS and context_dir.startswith("/"):
+    if context_dir not in CONTEXT_DIRS and is_local_path(context_dir):
         atexit.register(lambda: shutil.rmtree(context_dir, ignore_errors=True))
         CONTEXT_DIRS.add(context_dir)
 
diff --git a/cubed/utils.py b/cubed/utils.py
index fb01220e..23c5215a 100644
--- a/cubed/utils.py
+++ b/cubed/utils.py
@@ -74,6 +74,11 @@ def join_path(dir_url: PathType, child_path: str) -> str:
     return urlunsplit(split_parts)
 
 
+def is_local_path(path: str):
+    """Determine if a path string is for the local filesystem."""
+    return urlsplit(path).scheme in ("", "file")
+
+
 def memory_repr(num: int) -> str:
     """Convert bytes to a human-readable string in decimal form.
     1 KB is 1,000 bytes, 1 MB is 1,000,000 bytes, and so on.


Test patch:
diff --git a/cubed/tests/test_utils.py b/cubed/tests/test_utils.py
index 81a1358a..afe00c2b 100644
--- a/cubed/tests/test_utils.py
+++ b/cubed/tests/test_utils.py
@@ -12,6 +12,7 @@ from cubed.utils import (
     block_id_to_offset,
     broadcast_trick,
     extract_stack_summaries,
+    is_local_path,
     join_path,
     map_nested,
     memory_repr,
@@ -69,6 +70,15 @@ def test_join_path():
     assert join_path("http://host/a path", "subpath") == "http://host/a%20path/subpath"
 
 
+def test_is_local_path():
+    assert is_local_path("relative_path/path")
+    assert is_local_path("/absolute_path/path")
+    assert is_local_path("file:relative_path/path")
+    assert is_local_path("file://absolute_path/path")
+    assert is_local_path("file:///absolute_path/path")
+    assert not is_local_path("s3://host/path")
+
+
 def test_memory_repr():
     assert memory_repr(0) == "0 bytes"
     assert memory_repr(1) == "1 bytes"
Repository: cubed-dev/cubed
Pull request #551 (2024-08-13 17:03:52)

Cast inputs to Cubed arrays in `apply_ufunc`


Patch:
diff --git a/cubed/core/gufunc.py b/cubed/core/gufunc.py
index bc85e8e8..35a4b4a7 100644
--- a/cubed/core/gufunc.py
+++ b/cubed/core/gufunc.py
@@ -68,7 +68,13 @@ def apply_gufunc(
     # Main code:
 
     # Cast all input arrays to cubed
-    # args = [asarray(a) for a in args]  # TODO: do we need to support casting?
+    # Use a spec if there is one. Note that all args have to have the same spec, and
+    # this will be checked later when constructing the plan (see check_array_specs).
+    from cubed.array_api.creation_functions import asarray
+
+    specs = [a.spec for a in args if hasattr(a, "spec")]
+    spec = specs[0] if len(specs) > 0 else None
+    args = [asarray(a, spec=spec) for a in args]
 
     if len(input_coredimss) != len(args):
         raise ValueError(


Test patch:
diff --git a/cubed/tests/test_gufunc.py b/cubed/tests/test_gufunc.py
index a06d2e65..deb7d583 100644
--- a/cubed/tests/test_gufunc.py
+++ b/cubed/tests/test_gufunc.py
@@ -38,6 +38,16 @@ def test_apply_gufunc_elemwise_01(spec):
     assert_equal(z, np.array([2, 4, 6]))
 
 
+def test_apply_gufunc_elemwise_01_non_cubed_input(spec):
+    def add(x, y):
+        return x + y
+
+    a = cubed.from_array(np.array([1, 2, 3]), chunks=3, spec=spec)
+    b = np.array([1, 2, 3])
+    z = apply_gufunc(add, "(),()->()", a, b, output_dtypes=a.dtype)
+    assert_equal(z, np.array([2, 4, 6]))
+
+
 def test_apply_gufunc_elemwise_loop(spec):
     def foo(x):
         assert x.shape in ((2,), (1,))
Repository: cubed-dev/cubed
Pull request #608 (2024-11-05 15:08:29)

Fix groupby tests running on tensorstore


Patch:
diff --git a/cubed/core/groupby.py b/cubed/core/groupby.py
index 8ec97045..84a491a6 100644
--- a/cubed/core/groupby.py
+++ b/cubed/core/groupby.py
@@ -127,7 +127,7 @@ def groupby_blockwise(
     axis=None,
     dtype=None,
     num_groups=None,
-    extra_func_kwargs=None,
+    **kwargs,
 ):
     """A blockwise operation that performs groupby aggregations.
 
@@ -148,8 +148,6 @@ def groupby_blockwise(
         Data type of output.
     num_groups: int
         The number of groups in the grouping array ``by``.
-    extra_func_kwargs: dict, optional
-        Extra keyword arguments to pass to ``func``.
     """
 
     if by.ndim != 1:
@@ -203,6 +201,7 @@ def groupby_blockwise(
         by_read_chunks=by_read_chunks,
         target_chunks=target_chunks,
         groups_per_chunk=groups_per_chunk,
+        **kwargs,
     )
 
 


Test patch:
diff --git a/.github/workflows/tensorstore-tests.yml b/.github/workflows/tensorstore-tests.yml
index 74bfce99..827ff1df 100644
--- a/.github/workflows/tensorstore-tests.yml
+++ b/.github/workflows/tensorstore-tests.yml
@@ -1,6 +1,7 @@
 name: TensorStore tests
 
 on:
+  pull_request:
   schedule:
     # Every weekday at 03:58 UTC, see https://crontab.guru/
     - cron: "58 3 * * 1-5"
diff --git a/cubed/tests/test_groupby.py b/cubed/tests/test_groupby.py
index ab047893..6ca34c94 100644
--- a/cubed/tests/test_groupby.py
+++ b/cubed/tests/test_groupby.py
@@ -103,6 +103,7 @@ def test_groupby_blockwise_axis0():
         axis=0,
         dtype=nxp.int64,
         num_groups=6,
+        groupby_dtype=nxp.int32,
     )
     assert_array_equal(
         c.compute(),
@@ -129,6 +130,7 @@ def test_groupby_blockwise_axis1():
         axis=1,
         dtype=nxp.int64,
         num_groups=6,
+        groupby_dtype=nxp.int32,
     )
     assert_array_equal(
         c.compute(),
@@ -142,7 +144,9 @@ def test_groupby_blockwise_axis1():
     )
 
 
-def _sum_reduction_func(arr, by, axis, start_group, num_groups):
+def _sum_reduction_func(arr, by, axis, start_group, num_groups, groupby_dtype):
     # change 'by' so it starts from 0 for each chunk
     by = by - start_group
-    return npg.aggregate(by, arr, func="sum", axis=axis, size=num_groups)
+    return npg.aggregate(
+        by, arr, func="sum", dtype=groupby_dtype, axis=axis, size=num_groups
+    )
Repository: cubed-dev/cubed
Pull request #679 (2025-01-21 11:54:37)

Add scalar support to `where`
Part of #678 


Patch:
diff --git a/cubed/array_api/searching_functions.py b/cubed/array_api/searching_functions.py
index 93131db8..eeec8411 100644
--- a/cubed/array_api/searching_functions.py
+++ b/cubed/array_api/searching_functions.py
@@ -88,5 +88,13 @@ def _searchsorted(x, y, side):
 
 
 def where(condition, x1, x2, /):
+    x1_is_scalar = isinstance(x1, (int, float, complex, bool))
+    x2_is_scalar = isinstance(x2, (int, float, complex, bool))
+    if x1_is_scalar and x2_is_scalar:
+        raise TypeError("At least one of x1 and x2 must be an array in where")
+    elif x1_is_scalar:
+        x1 = x2._promote_scalar(x1)
+    elif x2_is_scalar:
+        x2 = x1._promote_scalar(x2)
     dtype = result_type(x1, x2)
     return elemwise(nxp.where, condition, x1, x2, dtype=dtype)


Test patch:
diff --git a/cubed/tests/test_array_api.py b/cubed/tests/test_array_api.py
index 6d2a41d1..a0b023b2 100644
--- a/cubed/tests/test_array_api.py
+++ b/cubed/tests/test_array_api.py
@@ -794,6 +794,22 @@ def test_searchsorted_sorter_not_implemented():
         xp.searchsorted(xp.asarray([1, 0]), xp.asarray([1]), sorter=xp.asarray([1, 0]))
 
 
+def test_where_scalars():
+    condition = xp.asarray(
+        [[True, False, True], [False, True, False], [True, False, True]], chunks=(2, 2)
+    )
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2))
+
+    b = xp.where(condition, a, 0)
+    assert_array_equal(b.compute(), np.array([[1, 0, 3], [0, 5, 0], [7, 0, 9]]))
+
+    c = xp.where(condition, 0, a)
+    assert_array_equal(c.compute(), np.array([[0, 2, 0], [4, 0, 6], [0, 8, 0]]))
+
+    with pytest.raises(TypeError):
+        xp.where(condition, 0, 1)
+
+
 # Statistical functions
Repository: cubed-dev/cubed
Pull request #693 (2025-02-11 13:09:15)

Support float32 in random number generation


Patch:
diff --git a/cubed/random.py b/cubed/random.py
index 6c60a6c9..4016f94e 100644
--- a/cubed/random.py
+++ b/cubed/random.py
@@ -9,14 +9,13 @@ from cubed.utils import block_id_to_offset, normalize_shape
 from cubed.vendor.dask.array.core import normalize_chunks
 
 
-def random(size, *, chunks=None, spec=None):
+def random(size, *, dtype=nxp.float64, chunks=None, spec=None):
     """Return random floats in the half-open interval [0.0, 1.0)."""
     shape = normalize_shape(size)
-    dtype = nxp.float64
     chunks = normalize_chunks(chunks, shape=shape, dtype=dtype)
     numblocks = tuple(map(len, chunks))
     root_seed = pyrandom.getrandbits(128)
-
+    extra_func_kwargs = dict(dtype=dtype)
     return map_blocks(
         _random,
         dtype=dtype,
@@ -24,12 +23,13 @@ def random(size, *, chunks=None, spec=None):
         spec=spec,
         numblocks=numblocks,
         root_seed=root_seed,
+        extra_func_kwargs=extra_func_kwargs,
     )
 
 
-def _random(x, numblocks=None, root_seed=None, block_id=None):
+def _random(x, numblocks=None, root_seed=None, dtype=nxp.float64, block_id=None):
     stream_id = block_id_to_offset(block_id, numblocks)
     rg = Generator(Philox(key=root_seed + stream_id))
-    out = rg.random(x.shape)
+    out = rg.random(x.shape, dtype=dtype)
     out = numpy_array_to_backend_array(out)
     return out


Test patch:
diff --git a/cubed/tests/test_random.py b/cubed/tests/test_random.py
index 5144538c..16ce795b 100644
--- a/cubed/tests/test_random.py
+++ b/cubed/tests/test_random.py
@@ -29,8 +29,22 @@ def test_random(spec, executor):
 
     assert a.shape == (10, 10)
     assert a.chunks == ((4, 4, 2), (5, 5))
+    assert a.dtype == xp.float64
 
     x = nxp.unique_values(a.compute(executor=executor))
+    assert x.dtype == xp.float64
+    assert len(x) > 90
+
+
+def test_random_dtype(spec, executor):
+    a = cubed.random.random((10, 10), dtype=xp.float32, chunks=(4, 5), spec=spec)
+
+    assert a.shape == (10, 10)
+    assert a.chunks == ((4, 4, 2), (5, 5))
+    assert a.dtype == xp.float32
+
+    x = nxp.unique_values(a.compute(executor=executor))
+    assert x.dtype == xp.float32
     assert len(x) > 90