Add linalg tests, modify squeeze for array api compliance

Author: lshaw8317

src/blosc2/lazyexpr.py

@@ -2900,7 +2900,7 @@ def __getitem__(self, item):
         # Squeeze single-element dimensions when indexing with integers
         # See e.g. examples/ndarray/animated_plot.py
         if isinstance(item, int) or (hasattr(item, "__iter__") and any(isinstance(i, int) for i in item)):
-            result = result.squeeze()
+            result = result.squeeze(axis=tuple(i for i in range(result.ndim) if result.shape[i] == 1))
         return result
 
     def slice(self, item):

src/blosc2/linalg.py

@@ -80,7 +80,7 @@ def matmul(x1: blosc2.Array, x2: blosc2.NDArray, **kwargs: Any) -> blosc2.NDArra
         raise ValueError("Arguments can't be scalars.")
 
     # Makes a SimpleProxy if inputs are not blosc2 arrays
-    x1, x2 = blosc2.as_simpleproxy(x1), blosc2.as_simpleproxy(x2)
+    x1, x2 = blosc2.as_simpleproxy(x1, x2)
 
     # Validate matrix multiplication compatibility
     if x1.shape[builtins.max(-1, -len(x2.shape))] != x2.shape[builtins.max(-2, -len(x2.shape))]:
@@ -184,7 +184,7 @@ def tensordot(
     # TODO: add fast path for when don't need to change chunkshapes
 
     # Makes a SimpleProxy if inputs are not blosc2 arrays
-    x1, x2 = blosc2.as_simpleproxy(x1), blosc2.as_simpleproxy(x2)
+    x1, x2 = blosc2.as_simpleproxy(x1, x2)
 
     if isinstance(axes, tuple):
         a_axes, b_axes = axes
@@ -326,7 +326,7 @@ def vecdot(x1: blosc2.NDArray, x2: blosc2.NDArray, axis: int = -1, **kwargs) ->
         return npvecdot(x1, x2, axis=axis)
 
     # Makes a SimpleProxy if inputs are not blosc2 arrays
-    x1, x2 = blosc2.as_simpleproxy(x1), blosc2.as_simpleproxy(x2)
+    x1, x2 = blosc2.as_simpleproxy(x1, x2)
 
     N = builtins.min(x1.ndim, x2.ndim)
     if axis < -N or axis > -1:
@@ -489,9 +489,15 @@ def permute_dims(
 
     chunks = arr.chunks
     shape = arr.shape
-
-    for info in arr.iterchunks_info():
-        coords = info.coords
+    # handle SimpleProxy which doesn't have iterchunks_info
+    if hasattr(arr, "iterchunks_info"):
+        my_it = arr.iterchunks_info()
+        _get_el = lambda x: x.coords  # noqa: E731
+    else:
+        my_it = get_intersecting_chunks((), shape, chunks)
+        _get_el = lambda x: x.raw  # noqa: E731
+    for info in my_it:
+        coords = _get_el(info)
         start_stop = [
             (coord * chunk, builtins.min(chunk * (coord + 1), dim))
             for coord, chunk, dim in zip(coords, chunks, shape, strict=False)
@@ -636,7 +642,7 @@ def outer(x1: blosc2.blosc2.NDArray, x2: blosc2.blosc2.NDArray, **kwargs: Any) -
     out: blosc2.NDArray
         A two-dimensional array containing the outer product and whose shape is (N, M).
     """
-    x1, x2 = blosc2.as_simpleproxy(x1), blosc2.as_simpleproxy(x2)
+    x1, x2 = blosc2.as_simpleproxy(x1, x2)
     if (x1.ndim != 1) or (x2.ndim != 1):
         raise ValueError("outer only valid for 1D inputs.")
     return tensordot(x1, x2, ((), ()), **kwargs)  # for testing purposes

src/blosc2/ndarray.py

@@ -4642,7 +4642,7 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray:
 
         return ndslice
 
-    def squeeze(self, axis=None) -> NDArray:
+    def squeeze(self, axis: int | Sequence[int]) -> 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.
@@ -4666,18 +4666,15 @@ def squeeze(self, axis=None) -> NDArray:
         >>> a.shape
         (23, 11)
         """
-        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)
+        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:
@@ -4711,17 +4708,23 @@ def __matmul__(self, other):
         return blosc2.linalg.matmul(self, other)
 
 
-def squeeze(x: NDArray, axis: int | None = None) -> NDArray:
+def squeeze(x: Array, axis: int | Sequence[int]) -> 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 axis is provided, it should be an int or tuple of ints and the corresponding
-    dimensions (of size 1) will be removed.
+    This method modifies the array in-place.
+
+    Parameters
+    ----------
+    x: Array
+        input array.
+    axis: int | Sequence[int]
+        Axis (or axes) to squeeze.
 
     Returns
     -------
-    out: NDArray
+    out: Array
+        An output array having the same data type and elements as x.
 
     Examples
     --------

src/blosc2/proxy.py

@@ -674,25 +674,28 @@ def __getitem__(self, item: slice | list[slice]) -> np.ndarray:
             return np.asarray(out)  # avoids copy for PyTorch at least
 
 
-def as_simpleproxy(x: blosc2.Array) -> SimpleProxy | blosc2.Operand:
+def as_simpleproxy(*arrs: Sequence[blosc2.Array]) -> tuple[SimpleProxy | blosc2.Operand]:
     """
     Convert an Array object which fulfills Array protocol into SimpleProxy. If x is already a
     blosc2.Operand simply returns object.
 
     Parameters
     ----------
-    x: blosc2.Array
-        Object fulfilling Array protocol.
+    arrs: Sequence[blosc2.Array]
+        Objects fulfilling Array protocol.
 
     Returns
     -------
-    out: blosc2.SimpleProxy | blosc2.Operand
-        Object with minimal interface for blosc2 LazyExpr computations.
+    out: tuple[blosc2.SimpleProxy | blosc2.Operand]
+        Objects with minimal interface for blosc2 LazyExpr computations.
     """
-    if isinstance(x, blosc2.Operand):
-        return x
-    else:
-        return SimpleProxy(x)
+    out = ()
+    for x in arrs:
+        if isinstance(x, blosc2.Operand):
+            out += (x,)
+        else:
+            out += (SimpleProxy(x),)
+    return out[0] if len(out) == 1 else out
 
 
 def jit(func=None, *, out=None, disable=False, **kwargs):  # noqa: C901

tests/ndarray/test_lazyexpr.py

@@ -10,6 +10,7 @@
 
 import numpy as np
 import pytest
+import torch
 
 import blosc2
 from blosc2.lazyexpr import ne_evaluate
@@ -1697,13 +1698,13 @@ def test_lazylinalg():
     np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- squeeze ---
-    out = blosc2.lazyexpr("squeeze(D)")
-    npres = np.squeeze(npD)
+    out = blosc2.lazyexpr("squeeze(D, axis=-1)")
+    npres = np.squeeze(npD, -1)
     assert out.shape == npres.shape
     np.testing.assert_array_almost_equal(out[()], npres)
 
-    out = blosc2.lazyexpr("D.squeeze()")
-    npres = np.squeeze(npD)
+    out = blosc2.lazyexpr("D.squeeze(axis=-1)")
+    npres = np.squeeze(npD, -1)
     assert out.shape == npres.shape
     np.testing.assert_array_almost_equal(out[()], npres)
 
@@ -1772,3 +1773,52 @@ def test_lazyexpr_2args():
     newexpr = blosc2.hypot(lexpr, 3)
     assert newexpr.expression == "hypot((sin(o0)), 3)"
     assert newexpr.operands["o0"] is a
+
+
+@pytest.mark.parametrize(
+    "xp",
+    [torch, np],
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["bool", "int32", "int64", "float32", "float64", "complex128"],
+)
+def test_simpleproxy(xp, dtype):
+    dtype_ = getattr(xp, dtype) if hasattr(xp, dtype) else np.dtype(dtype)
+    if dtype == "bool":
+        blosc_matrix = blosc2.asarray([True, False, False], dtype=np.dtype(dtype), chunks=(2,))
+        foreign_matrix = xp.zeros((3,), dtype=dtype_)
+        # Create a lazy expression object
+        lexpr = blosc2.lazyexpr(
+            "(b & a) | (~b)", operands={"a": blosc_matrix, "b": foreign_matrix}
+        )  # this does not
+        # Compare with numpy computation result
+        npb = np.asarray(foreign_matrix)
+        npa = blosc_matrix[()]
+        res = (npb & npa) | np.logical_not(npb)
+    else:
+        N = 10
+        shape_a = (N, N, N)
+        blosc_matrix = blosc2.full(shape=shape_a, fill_value=3, dtype=np.dtype(dtype), chunks=(N // 3,) * 3)
+        foreign_matrix = xp.ones(shape_a, dtype=dtype_)
+        if dtype == "complex128":
+            foreign_matrix += 0.5j
+            blosc_matrix = blosc2.full(
+                shape=shape_a, fill_value=3 + 2j, dtype=np.dtype(dtype), chunks=(N // 3,) * 3
+            )
+
+        # Create a lazy expression object
+        lexpr = blosc2.lazyexpr(
+            "b + sin(a) + sum(b) - tensordot(a, b, axes=1)",
+            operands={"a": blosc_matrix, "b": foreign_matrix},
+        )  # this does not
+        # Compare with numpy computation result
+        npb = np.asarray(foreign_matrix)
+        npa = blosc_matrix[()]
+        res = npb + np.sin(npa) + np.sum(npb) - np.tensordot(npa, npb, axes=1)
+
+    # Test object metadata and result
+    assert isinstance(lexpr, blosc2.LazyExpr)
+    assert lexpr.dtype == res.dtype
+    assert lexpr.shape == res.shape
+    np.testing.assert_array_equal(lexpr[()], res)

tests/ndarray/test_linalg.py

@@ -1,9 +1,12 @@
+import inspect
 from itertools import permutations
 
 import numpy as np
 import pytest
+import torch
 
 import blosc2
+from blosc2.lazyexpr import linalg_funcs
 from blosc2.ndarray import npvecdot
 
 
@@ -817,3 +820,54 @@ def test_diagonal(shape, chunkshape, offset):
 
     # Assert equality
     np.testing.assert_array_equal(result_np, expected)
+
+
+@pytest.mark.parametrize(
+    "xp",
+    [torch, np],
+)
+@pytest.mark.parametrize(
+    "dtype",
+    ["bool", "int32", "int64", "float32", "float64", "complex128"],
+)
+def test_linalgproxy(xp, dtype):
+    dtype_ = getattr(xp, dtype) if hasattr(xp, dtype) else np.dtype(dtype)
+    for name in linalg_funcs:
+        if name == "transpose":
+            continue  # deprecated
+        func = getattr(blosc2, name)
+        N = 10
+        shape_a = (N,)
+        chunks = (N // 3,)
+        if name != "outer":
+            shape_a *= 3
+            chunks *= 3
+        blosc_matrix = blosc2.full(shape=shape_a, fill_value=3, dtype=np.dtype(dtype), chunks=chunks)
+        foreign_matrix = xp.ones(shape_a, dtype=dtype_)
+        if dtype == "complex128":
+            foreign_matrix += 0.5j
+            blosc_matrix = blosc2.full(
+                shape=shape_a, fill_value=3 + 2j, dtype=np.dtype(dtype), chunks=chunks
+            )
+
+        # Check this works
+        argspec = inspect.getfullargspec(func)
+        num_args = len(argspec.args)
+        npfunc = blosc2.linalg.nptranspose if name == "permute_dims" else getattr(np, name)
+        if num_args > 2 or name in ("outer", "matmul"):
+            try:
+                lexpr = func(blosc_matrix, foreign_matrix)
+            except NotImplementedError:
+                continue
+            foreign_matrix = np.asarray(foreign_matrix)
+            res = npfunc(blosc_matrix[()], foreign_matrix)
+        else:
+            try:
+                lexpr = func(foreign_matrix)
+            except NotImplementedError:
+                continue
+            except TypeError:
+                continue
+            foreign_matrix = np.asarray(foreign_matrix)
+            res = npfunc(foreign_matrix, 0) if name == "expand_dims" else npfunc(foreign_matrix)
+        np.testing.assert_array_equal(res, lexpr[()])

tests/ndarray/test_proxy_expr.py

@@ -9,7 +9,6 @@
 
 import numpy as np
 import pytest
-import torch
 
 import blosc2
 from blosc2.lazyexpr import ne_evaluate
@@ -88,52 +87,3 @@ def test_expr_proxy_operands(chunks_blocks, c2sub_context):
     blosc2.remove_urlpath(urlpath)
     for path in cleanup_paths:
         blosc2.remove_urlpath(path)
-
-
-@pytest.mark.parametrize(
-    "xp",
-    [torch, np],
-)
-@pytest.mark.parametrize(
-    "dtype",
-    ["bool", "int32", "int64", "float32", "float64", "complex128"],
-)
-def test_simpleproxy(xp, dtype):
-    dtype_ = getattr(xp, dtype) if hasattr(xp, dtype) else np.dtype(dtype)
-    if dtype == "bool":
-        blosc_matrix = blosc2.asarray([True, False, False], dtype=np.dtype(dtype), chunks=(2,))
-        foreign_matrix = xp.zeros((3,), dtype=dtype_)
-        # Create a lazy expression object
-        lexpr = blosc2.lazyexpr(
-            "(b & a) | (~b)", operands={"a": blosc_matrix, "b": foreign_matrix}
-        )  # this does not
-        # Compare with numpy computation result
-        npb = np.asarray(foreign_matrix)
-        npa = blosc_matrix[()]
-        res = (npb & npa) | np.logical_not(npb)
-    else:
-        N = 10
-        shape_a = (N, N, N)
-        blosc_matrix = blosc2.full(shape=shape_a, fill_value=3, dtype=np.dtype(dtype), chunks=(N // 3,) * 3)
-        foreign_matrix = xp.ones(shape_a, dtype=dtype_)
-        if dtype == "complex128":
-            foreign_matrix += 0.5j
-            blosc_matrix = blosc2.full(
-                shape=shape_a, fill_value=3 + 2j, dtype=np.dtype(dtype), chunks=(N // 3,) * 3
-            )
-
-        # Create a lazy expression object
-        lexpr = blosc2.lazyexpr(
-            "b + sin(a) + sum(b) - tensordot(a, b, axes=1)",
-            operands={"a": blosc_matrix, "b": foreign_matrix},
-        )  # this does not
-        # Compare with numpy computation result
-        npb = np.asarray(foreign_matrix)
-        npa = blosc_matrix[()]
-        res = npb + np.sin(npa) + np.sum(npb) - np.tensordot(npa, npb, axes=1)
-
-    # Test object metadata and result
-    assert isinstance(lexpr, blosc2.LazyExpr)
-    assert lexpr.dtype == res.dtype
-    assert lexpr.shape == res.shape
-    np.testing.assert_array_equal(lexpr[()], res)