Add extract operation with tests

Author: jim22k

mlir_graphblas/implementations.py

@@ -900,3 +900,45 @@ def main(x):
             main.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
 
         return compile(module)
+
+
+def extract(tensor: SparseTensorBase, row_indices, col_indices=None):
+    # There may be a way to do this in MLIR, but for now we use numpy
+    if tensor.ndims == 1:
+        # Vector
+        assert col_indices is None
+        if row_indices is None:  # None indicate GrB_ALL
+            return tensor.dup()
+        rowidx, vals = tensor.extract_tuples()
+        selected = np.isin(rowidx, row_indices)
+        v = Vector.new(tensor.dtype, *tensor.shape)
+        v.build(rowidx[selected], vals[selected])
+        return v
+
+    # Matrix
+    if row_indices is None and col_indices is None:
+        return tensor.dup()
+    rowidx, colidx, vals = tensor.extract_tuples()
+    if row_indices is not None:
+        rowsel = np.isin(rowidx, row_indices)
+        # Apply rowsel filter
+        rowidx, colidx, vals = rowidx[rowsel], colidx[rowsel], vals[rowsel]
+    if col_indices is not None:
+        colsel = np.isin(colidx, col_indices)
+        # Apply colsel filter
+        rowidx, colidx, vals = rowidx[colsel], colidx[colsel], vals[colsel]
+    if type(row_indices) is int:
+        # Extract row as Vector
+        assert np.all(rowidx == row_indices)
+        v = Vector.new(tensor.dtype, tensor.shape[1])
+        v.build(colidx, vals)
+        return v
+    if type(col_indices) is int:
+        # Extract col as Vector
+        assert np.all(colidx == col_indices)
+        v = Vector.new(tensor.dtype, tensor.shape[0])
+        v.build(rowidx, vals)
+        return v
+    m = Matrix.new(tensor.dtype, *tensor.shape)
+    m.build(rowidx, colidx, vals)
+    return m

mlir_graphblas/operations.py

@@ -1,3 +1,4 @@
+import numpy as np
 from typing import Optional, Union
 from .tensor import SparseObject, SparseTensor, SparseTensorBase, Matrix, Vector, Scalar, TransposedMatrix
 from .operators import UnaryOp, BinaryOp, SelectOp, IndexUnaryOp, Monoid, Semiring
@@ -479,3 +480,58 @@ def reduce_to_scalar(out: Scalar,
 
     result = impl.reduce_to_scalar(op, tensor)
     update(out, result, accum=accum, desc=desc)
+
+
+def extract(out: SparseTensor,
+            tensor: SparseTensorBase,
+            row_indices=None,
+            col_indices=None,
+            *,
+            mask: Optional[Vector] = None,
+            accum: Optional[BinaryOp] = None,
+            desc: Descriptor = NULL_DESC):
+    """
+    Setting row_indices or col_indices to `None` is the equivalent of GrB_ALL
+    """
+    # Verify dtypes
+    if out.dtype != tensor.dtype:
+        raise GrbDomainMismatch(f"output must have same dtype as input: {out.dtype} != {tensor.dtype}")
+
+    # Apply transpose
+    if desc.transpose0 and tensor.ndims == 2:
+        tensor = TransposedMatrix.wrap(tensor)
+
+    # Check indices
+    if tensor.ndims < 1:
+        raise TypeError("Use `extract_element` rather than `extract` for Scalars")
+    if tensor.ndims < 2 and col_indices is not None:
+        raise ValueError("col_indices not allowed for Vector, use row_indices")
+
+    # Compare shapes
+    if type(row_indices) is int and type(col_indices) is int:
+        raise TypeError("Cannot provide int for both row_indices and col_indices")
+    if type(row_indices) is int:
+        expected_out_shape = (tensor.shape[1],)
+    elif type(col_indices) is int:
+        expected_out_shape = (tensor.shape[0],)
+    else:
+        expected_out_shape = tensor.shape
+    if out.shape != expected_out_shape:
+        raise GrbDimensionMismatch(f"output shape mismatch: {out.shape} != {expected_out_shape}")
+
+    if mask is not None:
+        tensor = impl.apply_mask(tensor, mask, desc)
+
+    result = impl.extract(tensor, row_indices, col_indices)
+    update(out, result, mask, accum, desc)
+
+
+def assign(out: SparseTensor,
+           tensor: SparseObject,
+           row_indices=None,
+           col_indices=None,
+           *,
+           mask: Optional[Vector] = None,
+           accum: Optional[BinaryOp] = None,
+           desc: Descriptor = NULL_DESC):
+    raise NotImplementedError()

mlir_graphblas/tensor.py

@@ -119,6 +119,9 @@ def clear(self):
         self._rtt = None
         self._intermediate_result = False
 
+    def extract_tuples(self):
+        raise NotImplementedError()
+
 
 class SparseTensor(SparseTensorBase):
     def clear(self):
@@ -459,3 +462,7 @@ def nvals(self):
     def clear(self):
         super().clear()
         self._referenced_matrix = None
+
+    def extract_tuples(self):
+        col_indices, row_indices, values = self._referenced_matrix.extract_tuples()
+        return row_indices, col_indices, values

mlir_graphblas/tests/test_operations.py

@@ -202,14 +202,14 @@ def test_select_vec(vs):
     operations.select(z, SelectOp.rowgt, x, 2)
     idx, vals = z.extract_tuples()
     np_assert_equal(idx, [3])
-    np_assert_equal(vals, [30.])
+    np_assert_allclose(vals, [30.])
 
     # Select by value
     z = Vector.new(x.dtype, x.size())
     operations.select(z, SelectOp.valuegt, x, 10.)
     idx, vals = z.extract_tuples()
     np_assert_equal(idx, [2, 3])
-    np_assert_equal(vals, [20., 30.])
+    np_assert_allclose(vals, [20., 30.])
 
 
 def test_select_mat(mm):
@@ -220,7 +220,7 @@ def test_select_mat(mm):
     rowidx, colidx, vals = z.extract_tuples()
     np_assert_equal(rowidx, [0, 1, 1, 2])
     np_assert_equal(colidx, [4, 0, 4, 3])
-    np_assert_equal(vals, [6., 1., 8., 2.])
+    np_assert_allclose(vals, [6., 1., 8., 2.])
 
     # Transposed
     z = Matrix.new(y.dtype, y.shape[1], y.shape[0])
@@ -229,7 +229,7 @@ def test_select_mat(mm):
     rowidx, colidx, vals = z.extract_tuples()
     np_assert_equal(rowidx, [0, 0, 0])
     np_assert_equal(colidx, [1, 3, 5])
-    np_assert_equal(vals, [1., 5., 7.])
+    np_assert_allclose(vals, [1., 5., 7.])
 
 
 def test_empty_select():
@@ -275,3 +275,61 @@ def test_reduce_scalar_vec(vs):
     s = Scalar.new(x.dtype)
     operations.reduce_to_scalar(s, Monoid.times, x)
     np_assert_allclose(s.extract_element(), functools.reduce(operator.mul, xvals))
+
+
+def test_extract_vec(vs):
+    x, _ = vs
+    xidx, xvals = x.extract_tuples()
+    z = Vector.new(x.dtype, *x.shape)
+    operations.extract(z, x, [0, 1, 3])
+    idx, vals = z.extract_tuples()
+    np_assert_equal(idx, [1, 3])
+    np_assert_allclose(vals, [10., 30.])
+
+    # None == GrB_ALL
+    operations.extract(z, x, None)
+    idx, vals = z.extract_tuples()
+    np_assert_equal(idx, xidx)
+    np_assert_allclose(vals, xvals)
+
+
+def test_extract_mat(mm):
+    x, _ = mm
+    xrows, xcols, xvals = x.extract_tuples()
+    z = Matrix.new(x.dtype, *x.shape)
+    operations.extract(z, x, [0, 4], [1, 3, 5])
+    rowidx, colidx, vals = z.extract_tuples()
+    np_assert_equal(rowidx, [0, 0])
+    np_assert_equal(colidx, [3, 5])
+    np_assert_allclose(vals, [1.1, 2.2])
+
+    # None == GrB_ALL
+    operations.extract(z, x, None, None)
+    rowidx, colidx, vals = z.extract_tuples()
+    np_assert_equal(rowidx, xrows)
+    np_assert_equal(colidx, xcols)
+    np_assert_allclose(vals, xvals)
+
+
+def test_extract_vec_from_mat(mm):
+    x, _ = mm
+    # Extract column
+    z = Vector.new(x.dtype, x.shape[0])
+    operations.extract(z, x, [0, 1, 4], 3)
+    idx, vals = z.extract_tuples()
+    np_assert_equal(idx, [0, 1])
+    np_assert_allclose(vals, [1.1, 3.3])
+
+    # Extract row
+    z = Vector.new(x.dtype, x.shape[1])
+    operations.extract(z, x, 2, [0, 1, 4])
+    idx, vals = z.extract_tuples()
+    np_assert_equal(idx, [0, 1])
+    np_assert_allclose(vals, [4.4, 5.5])
+
+    # Extract column via transposed input
+    z = Vector.new(x.dtype, x.shape[0])
+    operations.extract(z, x, 3, [0, 1, 4], desc=desc.T0)
+    idx, vals = z.extract_tuples()
+    np_assert_equal(idx, [0, 1])
+    np_assert_allclose(vals, [1.1, 3.3])