diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index a1e18152b..1ed02aed5 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -735,6 +735,7 @@ def wait(self, how="materialize"): else: raise ValueError(f'`how` argument must be "materialize" or "complete"; got {how!r}') call("GrB_Matrix_wait", [self, mode]) + return self def get(self, row, col, default=None): """Get an element at (``row``, ``col``) indices as a Python scalar. diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index cc34e27e2..56064534d 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -465,6 +465,7 @@ def wait(self, how="materialize"): raise ValueError(f'`how` argument must be "materialize" or "complete"; got {how!r}') if not self._is_cscalar: call("GrB_Scalar_wait", [self, mode]) + return self def get(self, default=None): """Get the internal value of the Scalar as a Python scalar. diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index a39a92f40..fa658b07b 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -624,6 +624,7 @@ def wait(self, how="materialize"): else: raise ValueError(f'`how` argument must be "materialize" or "complete"; got {how!r}') call("GrB_Vector_wait", [self, mode]) + return self def get(self, index, default=None): """Get an element at ``index`` as a Python scalar. diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 351e44e64..6d58d82c5 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -4234,6 +4234,12 @@ def test_ss_descriptors(A): (A @ A).new(nthreads=4, axb_method="dot", sort=True) +@autocompute +def test_wait_chains(A): + result = A.wait().T.wait().reduce_rowwise().wait().reduce().wait() + assert result == 47 + + def test_subarray_dtypes(): a = np.arange(3 * 4, dtype=np.int64).reshape(3, 4) A = Matrix.from_coo([1, 3, 5], [0, 1, 3], a)