Add matmul support for N-D arrays, modify squeeze to comply with array-api

Author: lshaw8317

src/blosc2/lazyexpr.py

@@ -40,7 +40,13 @@
 import blosc2
 from blosc2 import compute_chunks_blocks
 from blosc2.info import InfoReporter
-from blosc2.ndarray import _check_allowed_dtypes, get_chunks_idx, is_inside_new_expr, process_key
+from blosc2.ndarray import (
+    _check_allowed_dtypes,
+    get_chunks_idx,
+    get_intersecting_chunks,
+    is_inside_new_expr,
+    process_key,
+)
 
 if not blosc2.IS_WASM:
     import numexpr
@@ -1575,7 +1581,7 @@ def slices_eval(  # noqa: C901
             elif isinstance(out, blosc2.NDArray):
                 # It *seems* better to choose an automatic chunks and blocks for the output array
                 # out = out.slice(_slice, chunks=out.chunks, blocks=out.blocks)
-                out = out.squeeze(mask_slice)
+                out = out.squeeze(np.where(mask_slice)[0])
             else:
                 raise ValueError("The output array is not a NumPy array or a NDArray")
 
@@ -3629,16 +3635,6 @@ def evaluate(
     return lexpr[()]
 
 
-def get_intersecting_chunks(_slice, shape, chunks):
-    if 0 not in chunks:
-        chunk_size = ndindex.ChunkSize(chunks)
-        return chunk_size.as_subchunks(_slice, shape)  # if _slice is (), returns all chunks
-    else:
-        return (
-            ndindex.ndindex(...).expand(shape),
-        )  # chunk is whole array so just return full tuple to do loop once
-
-
 if __name__ == "__main__":
     from time import time
 

src/blosc2/ndarray.py

@@ -2288,7 +2288,7 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray:
                 for order, nchunk in enumerate(aligned_chunks):
                     chunk = self.schunk.get_chunk(nchunk)
                     newarr.schunk.update_chunk(order, chunk)
-                newarr.squeeze(mask=mask)  # remove any dummy dims introduced
+                newarr.squeeze(axis=np.where(mask)[0])  # remove any dummy dims introduced
                 return newarr
 
         key = (start, stop)
@@ -2307,11 +2307,11 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray:
 
         return ndslice
 
-    def squeeze(self, mask=None) -> NDArray:
+    def squeeze(self, axis=None) -> NDArray:
         """Remove single-dimensional entries from the shape of the array.
 
         This method modifies the array in-place. If mask is None removes any dimensions with size 1.
-        If mask is provided, it should be a boolean array of the same shape as the array, and the corresponding
+        If axis is provided, it should be an int or tuple of ints and the corresponding
         dimensions (of size 1) will be removed.
 
         Returns
@@ -2331,7 +2331,18 @@ def squeeze(self, mask=None) -> NDArray:
         >>> a.shape
         (23, 11)
         """
-        super().squeeze(mask=mask)
+        if axis is None:
+            super().squeeze()
+        else:
+            axis = [axis] if isinstance(axis, int) else axis
+            mask = [False for i in range(self.ndim)]
+            for a in axis:
+                if a < 0:
+                    a += self.ndim  # Adjust axis to be within the array's dimensions
+                if mask[a]:
+                    raise ValueError("Axis values must be unique.")
+                mask[a] = True
+            super().squeeze(mask=mask)
         return self
 
     def indices(self, order: str | list[str] | None = None, **kwargs: Any) -> NDArray:
@@ -4312,9 +4323,8 @@ def asarray(
     else:
         if not isinstance(array, NDArray):
             raise ValueError("Must always do a copy for asarray unless NDArray provided.")
-        mask = [True] + [False for i in range(array.ndim)]
         # TODO: make a direct view possible
-        return blosc2.expand_dims(array, axis=0).squeeze(mask)  # way to get a view
+        return blosc2.expand_dims(array, axis=0).squeeze(axis=0)  # way to get a view
 
     return ndarr
 
@@ -4515,15 +4525,15 @@ def sort(array: NDArray, order: str | list[str] | None = None, **kwargs: Any) ->
     return larr.sort(order).compute(**kwargs)
 
 
-def matmul(x1: NDArray, x2: NDArray, **kwargs: Any) -> NDArray:
+def matmul(x1: NDArray | np.ndarray, x2: NDArray, **kwargs: Any) -> NDArray | np.ndarray:  # noqa : C901
     """
     Computes the matrix product between two Blosc2 NDArrays.
 
     Parameters
     ----------
-    x1: :ref:`NDArray`
+    x1: :ref:`NDArray` | np.ndarray
         The first input array.
-    x2: :ref:`NDArray`
+    x2: :ref:`NDArray` | np.ndarray
         The second input array.
     kwargs: Any, optional
         Keyword arguments that are supported by the :func:`empty` constructor.
@@ -4575,51 +4585,70 @@ def matmul(x1: NDArray, x2: NDArray, **kwargs: Any) -> NDArray:
     array([1, 5])
 
     """
+    # Added this to pass array-api tests (which use internal getitem to check results)
+    if isinstance(x1, np.ndarray) and isinstance(x2, np.ndarray):
+        return np.matmul(x1, x2)
 
     # Validate arguments are not scalars
     if np.isscalar(x1) or np.isscalar(x2):
         raise ValueError("Arguments can't be scalars.")
 
-    # Validate arguments are dimension 1 or 2
-    if x1.ndim > 2 or x2.ndim > 2:
-        raise ValueError("Multiplication of arrays with dimension greater than 2 is not supported yet.")
+    # Validate matrix multiplication compatibility
+    if x1.shape[-1] != x2.shape[builtins.max(-2, -len(x2.shape))]:
+        raise ValueError("Shapes are not aligned for matrix multiplication.")
 
     # Promote 1D arrays to 2D if necessary
     x1_is_vector = False
     x2_is_vector = False
     if x1.ndim == 1:
-        x1 = x1.reshape((1, x1.shape[0]))  # (N,) -> (1, N)
+        x1 = blosc2.expand_dims(x1, axis=0)  # (N,) -> (1, N)
         x1_is_vector = True
     if x2.ndim == 1:
-        x2 = x2.reshape((x2.shape[0], 1))  # (M,) -> (M, 1)
+        x2 = blosc2.expand_dims(x2, axis=1)  # (M,) -> (M, 1)
         x2_is_vector = True
 
-    # Validate matrix multiplication compatibility
-    if x1.shape[-1] != x2.shape[-2]:
-        raise ValueError("Shapes are not aligned for matrix multiplication.")
-
     n, k = x1.shape[-2:]
     m = x2.shape[-1]
+    result_shape = np.broadcast_shapes(x1.shape[:-2], x2.shape[:-2]) + (n, m)
+    result = blosc2.zeros(result_shape, dtype=np.result_type(x1, x2), **kwargs)
 
-    result = blosc2.zeros((n, m), dtype=np.result_type(x1, x2), **kwargs)
+    if 0 in result.shape + x1.shape + x2.shape:  # if any array is empty, return array of 0s
+        if x1_is_vector:
+            result.squeeze(axis=-2)
+        if x2_is_vector:
+            result.squeeze(axis=-1)
+        return result
 
     p, q = result.chunks[-2:]
     r = x2.chunks[-1]
 
-    for row in range(0, n, p):
-        row_end = builtins.min(row + p, n)
-        for col in range(0, m, q):
-            col_end = builtins.min(col + q, m)
-            for aux in range(0, k, r):
-                aux_end = builtins.min(aux + r, k)
-                bx1 = x1[row:row_end, aux:aux_end]
-                bx2 = x2[aux:aux_end, col:col_end]
-                result[row:row_end, col:col_end] += np.matmul(bx1, bx2)
+    intersecting_chunks = get_intersecting_chunks((), result.shape[:-2], result.chunks[:-2])
+    for chunk in intersecting_chunks:
+        chunk = chunk.raw
+        for row in range(0, n, p):
+            row_end = builtins.min(row + p, n)
+            for col in range(0, m, q):
+                col_end = builtins.min(col + q, m)
+                for aux in range(0, k, r):
+                    aux_end = builtins.min(aux + r, k)
+                    bx1 = (
+                        x1[chunk[-x1.ndim + 2 :] + (slice(row, row_end), slice(aux, aux_end))]
+                        if x1.ndim > 2
+                        else x1[row:row_end, aux:aux_end]
+                    )
+                    bx2 = (
+                        x2[chunk[-x2.ndim + 2 :] + (slice(aux, aux_end), slice(col, col_end))]
+                        if x2.ndim > 2
+                        else x2[aux:aux_end, col:col_end]
+                    )
+                    result[chunk + (slice(row, row_end), slice(col, col_end))] += np.matmul(bx1, bx2)
 
-    if x1_is_vector and x2_is_vector:
-        return result[0][0]
+    if x1_is_vector:
+        result.squeeze(axis=-2)
+    if x2_is_vector:
+        result.squeeze(axis=-1)
 
-    return result.squeeze()
+    return result
 
 
 def permute_dims(
@@ -5178,6 +5207,16 @@ def _get_local_slice(prior_selection, post_selection, chunk_bounds):
     return locbegin, locend
 
 
+def get_intersecting_chunks(_slice, shape, chunks):
+    if 0 not in chunks:
+        chunk_size = ndindex.ChunkSize(chunks)
+        return chunk_size.as_subchunks(_slice, shape)  # if _slice is (), returns all chunks
+    else:
+        return (
+            ndindex.ndindex(...).expand(shape),
+        )  # chunk is whole array so just return full tuple to do loop once
+
+
 def broadcast_to(arr, shape):
     """
     Broadcast an array to a new shape.

tests/ndarray/test_matmul.py

@@ -9,6 +9,8 @@
     {
         ((12, 10), (7, 5), (3, 3)),
         ((10,), (9,), (7,)),
+        ((0,), (0,), (0,)),
+        ((40, 10, 10), (2, 3, 4), (1, 2, 2)),
     },
 )
 @pytest.mark.parametrize(
@@ -17,6 +19,9 @@
         ((10,), (4,), (2,)),
         ((10, 5), (3, 4), (1, 3)),
         ((10, 12), (2, 4), (1, 2)),
+        ((200, 10, 22), (23, 2, 4), (4, 1, 2)),
+        ((0,), (0,), (0,)),
+        ((20, 40, 10, 10), (5, 2, 3, 4), (2, 1, 2, 2)),
     },
 )
 @pytest.mark.parametrize(
@@ -26,13 +31,21 @@
 def test_matmul(ashape, achunks, ablocks, bshape, bchunks, bblocks, dtype):
     a = blosc2.linspace(0, 1, dtype=dtype, shape=ashape, chunks=achunks, blocks=ablocks)
     b = blosc2.linspace(0, 1, dtype=dtype, shape=bshape, chunks=bchunks, blocks=bblocks)
-    c = blosc2.matmul(a, b)
-
-    na = a[:]
-    nb = b[:]
-    nc = np.matmul(na, nb)
-
-    np.testing.assert_allclose(c, nc, rtol=1e-6)
+    a_np = a[:]
+    b_np = b[:]
+    try:
+        np_res = np.matmul(a_np, b_np)
+        np_error = None
+    except ValueError as e:
+        np_res = None
+        np_error = e
+
+    if np_error is not None:
+        with pytest.raises(type(np_error)):
+            blosc2.matmul(a, b)
+    else:
+        b2_res = blosc2.matmul(a, b)
+        np.testing.assert_allclose(b2_res[()], np_res, rtol=1e-6)
 
 
 @pytest.mark.parametrize(
@@ -147,12 +160,22 @@ def test_scalars(scalar):
 def test_dims(ashape, bshape):
     a = blosc2.linspace(0, 10, shape=ashape)
     b = blosc2.linspace(0, 1, shape=bshape)
-
-    with pytest.raises(ValueError):
-        blosc2.matmul(a, b)
-
-    with pytest.raises(ValueError):
-        blosc2.matmul(b, a)
+    a_np = a[:]
+    b_np = b[:]
+
+    try:
+        np_res = np.matmul(a_np, b_np)
+        np_error = None
+    except ValueError as e:
+        np_res = None
+        np_error = e
+
+    if np_error is not None:
+        with pytest.raises(type(np_error)):
+            blosc2.matmul(a, b)
+    else:
+        b2_res = blosc2.matmul(a, b)
+        np.testing.assert_allclose(b2_res[:], np_res)
 
 
 @pytest.mark.parametrize(

tests/ndarray/test_squeeze.py

@@ -13,19 +13,20 @@
 
 
 @pytest.mark.parametrize(
-    ("shape", "chunks", "blocks", "fill_value", "mask"),
+    ("shape", "chunks", "blocks", "fill_value", "axis"),
     [
-        ((1, 1230), (1, 100), (1, 3), b"0123", [True, False]),
-        ((23, 1, 1, 34), (20, 1, 1, 20), None, 1234, [False, False, True, False]),
-        ((80, 1, 51, 60, 1), None, (6, 1, 6, 26, 1), 3.333, [False] * 4 + [True]),
-        ((1, 1, 1), None, None, True, [False, True, True]),
+        ((1, 1230), (1, 100), (1, 3), b"0123", 0),
+        ((23, 1, 1, 34), (20, 1, 1, 20), None, 1234, 2),
+        ((80, 1, 51, 60, 1), None, (6, 1, 6, 26, 1), 3.333, 4),
+        ((1, 1, 1), None, None, True, (1, 2)),
+        ((1, 1, 1), None, None, True, None),
     ],
 )
-def test_squeeze(shape, chunks, blocks, fill_value, mask):
+def test_squeeze(shape, chunks, blocks, fill_value, axis):
     a = blosc2.full(shape, fill_value=fill_value, chunks=chunks, blocks=blocks)
 
-    b = np.squeeze(a[...], tuple(i for i, m in enumerate(mask) if m))
-    a_ = a.squeeze(mask)
+    b = np.squeeze(a[...], axis)
+    a_ = a.squeeze(axis)
 
     assert a_.shape == b.shape
     # TODO: this would work if squeeze returns a view