Add shape parsing of linalg funcs

Author: lshaw8317

src/blosc2/lazyexpr.py

@@ -620,6 +620,11 @@ def compute_smaller_slice(larger_shape, smaller_shape, larger_slice):
 valid_methods |= {"int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"}
 valid_methods |= {"float32", "float64", "complex64", "complex128"}
 valid_methods |= {"bool", "str", "bytes"}
+valid_methods |= {
+    name
+    for name in dir(blosc2.NDArray)
+    if callable(getattr(blosc2.NDArray, name)) and not name.startswith("_")
+}
 
 
 def validate_expr(expr: str) -> None:
@@ -2002,7 +2007,7 @@ def reduce_slices(  # noqa: C901
             continue
 
         if where is None:
-            if expression == "o0":
+            if expression == "o0" or expression == "(o0)":
                 # We don't have an actual expression, so avoid a copy except to make contiguous
                 result = np.require(chunk_operands["o0"], requirements="C")
             else:
@@ -3168,9 +3173,6 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
             # in guessing mode to avoid computing reductions
             # Extract possible numpy scalars
             _expression, local_vars = extract_numpy_scalars(expression)
-            # Let's include numpy and blosc2 as operands so that some functions can be used
-            # Most in particular, castings like np.int8 et al. can be very useful to allow
-            # for desired data types in the output.
             _operands = operands | local_vars
             # Check that operands are proper Operands, LazyArray or scalars; if not, convert to NDArray objects
             for op, val in _operands.items():
@@ -3179,10 +3181,10 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
             # for scalars just return value (internally converts to () if necessary)
             opshapes = {k: v if not hasattr(v, "shape") else v.shape for k, v in _operands.items()}
             _shape = infer_shape(_expression, opshapes)  # infer shape, includes constructors
-            # substitutes with numpy operands (cheap for reductions) and
-            # defaults to blosc2 functions (cheap for constructors)
             # have to handle slices since a[10] on a dummy variable of shape (1,1) doesn't work
             desliced_expr, desliced_ops = extract_and_replace_slices(_expression, _operands)
+            # substitutes with dummy operands (cheap for reductions) and
+            # defaults to blosc2 functions (cheap for constructors)
             new_expr = _numpy_eval_expr(desliced_expr, desliced_ops, prefer_blosc=True)
             _dtype = new_expr.dtype
             if isinstance(new_expr, blosc2.LazyExpr):
@@ -3205,24 +3207,16 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
                         if counter == 0 and char == ",":
                             break
                     expression_ = finalexpr[:-1]  # remove trailing comma
-                new_expr.expression = f"({expression_})"  # force parenthesis
-                new_expr.expression_tosave = expression
-                new_expr.operands = operands_
-                new_expr.operands_tosave = operands
-            elif isinstance(new_expr, blosc2.NDArray) and len(operands) == 1:
-                # passed "a", "a[:10]", 'sum(a)'
-                expression_, operands_ = conserve_functions(
-                    _expression, _operands, {"o0": list(operands.values())[0]} | local_vars
-                )
-                new_expr = cls(None)
-                new_expr.expression = expression_
-                new_expr.operands = operands_
             else:
+                new_expr = cls(None)
                 # An immediate evaluation happened
                 # (e.g. all operands are numpy arrays or constructors)
-                new_expr = cls(None)
-                new_expr.expression = expression
-                new_expr.operands = operands
+                # or passed "a", "a[:10]", 'sum(a)'
+                expression_, operands_ = conserve_functions(_expression, _operands, local_vars)
+            new_expr.expression = f"({expression_})"  # force parenthesis
+            new_expr.operands = operands_
+            new_expr.expression_tosave = expression
+            new_expr.operands_tosave = operands
             # Cache the dtype and shape (should be immutable)
             new_expr._dtype = _dtype
             new_expr._shape = _shape

src/blosc2/linalg.py

@@ -353,10 +353,6 @@ def vecdot(x1: blosc2.NDArray, x2: blosc2.NDArray, axis: int = -1, **kwargs) ->
     a_keep[a_axes] = False
     b_keep = [True] * x2.ndim
     b_keep[b_axes] = False
-    x1shape = np.array(x1.shape)
-    x2shape = np.array(x2.shape)
-    result_shape = np.broadcast_shapes(x1shape[a_keep], x2shape[b_keep])
-    result = blosc2.zeros(result_shape, dtype=np.result_type(x1, x2), **kwargs)
 
     x1shape = np.array(x1.shape)
     x2shape = np.array(x2.shape)

src/blosc2/shape_utils.py

@@ -1,4 +1,5 @@
 import ast
+import builtins
 
 from numpy import broadcast_shapes
 
@@ -15,6 +16,8 @@
     "tensordot",
     "transpose",
     "vecdot",
+    "T",
+    "mT",
 )
 reducers = ("sum", "prod", "min", "max", "std", "mean", "var", "any", "all", "slice", "count_nonzero")
 
@@ -39,6 +42,152 @@
 
 
 # --- Shape utilities ---
+def linalg_shape(func_name, args, kwargs):  # noqa: C901
+    # --- Linear algebra and tensor manipulation ---
+    a = args[0] if args else None
+    if a is None or any(s is None for s in a):
+        return None
+    b = args[1] if len(args) > 1 else None
+    axis = kwargs.get("axis", None)
+    axes = kwargs.get("axes", None)
+    offset = kwargs.get("offset", 0)
+
+    # --- concat ---
+    if func_name == "concat":
+        shapes = args[0]
+        if axis is None and len(args) > 1:
+            axis = args[1]
+
+        # Coerce axis to int if tuple single-element
+        axis = 0 if axis is None else axis
+        # normalize negative axis
+        axis = axis + len(shapes[0]) if axis < 0 else axis
+        concat_dim = builtins.sum([s[axis] for s in shapes])
+        return tuple(s if i != axis else concat_dim for i, s in enumerate(shapes[0]))
+
+    # --- diagonal ---
+    elif func_name == "diagonal":
+        axis1 = len(a) - 2
+        axis2 = len(a) - 1
+        new_shape = [d for i, d in enumerate(a) if i not in (axis1, axis2)]
+        d1, d2 = a[axis1], a[axis2]
+        diag_len = builtins.max(0, min(d1, d2) - abs(offset))
+        new_shape.append(diag_len)
+        return tuple(new_shape)
+
+    # --- expand_dims ---
+    elif func_name == "expand_dims":
+        # positional axis may be second positional argument
+        if axis is None and len(args) > 1:
+            axis = args[1]
+        if axis is None:
+            axis = 0
+        axis = [axis] if isinstance(axis, int) else axis
+        new_shape = list(a)
+        for ax in sorted(axis):
+            ax = ax if ax >= 0 else len(new_shape) + ax + 1
+            new_shape.insert(ax, 1)
+        return tuple(new_shape)
+
+    # --- matmul ---
+    elif func_name == "matmul":
+        if b is None:
+            return None
+        x1_is_vector = False
+        x2_is_vector = False
+        if len(a) == 1:
+            a = (1,) + a  # (N,) -> (1, N)
+            x1_is_vector = True
+        if len(b) == 1:
+            b += (1,)  # (M,) -> (M, 1)
+            x2_is_vector = True
+        batch = broadcast_shapes(a[:-2], b[:-2])
+        shape = batch
+        if not x1_is_vector:
+            shape += a[-2]
+        if not x2_is_vector:
+            shape += b[-1]
+        return shape
+
+    # --- matrix_transpose ---
+    elif func_name == "matrix_transpose":
+        if len(a) < 2:
+            return a
+        return a[:-2] + (a[-1], a[-2])
+
+    # --- outer ---
+    elif func_name == "outer":
+        if b is None:
+            return None
+        return a + b
+
+    # --- permute_dims ---
+    elif func_name == "permute_dims":
+        if axes is None and len(args) > 1:
+            axes = args[1]
+        if axes is None:
+            axes = tuple(reversed(range(len(a))))
+        return tuple(a[i] for i in axes)
+
+    # --- squeeze ---
+    elif func_name == "squeeze":
+        if axis is None and len(args) > 1:
+            axis = args[1]
+        if axis is None:
+            return tuple(d for d in a if d != 1)
+        if isinstance(axis, int):
+            axis = (axis,)
+        axis = tuple(ax if ax >= 0 else len(a) + ax for ax in axis)
+        return tuple(d for i, d in enumerate(a) if i not in axis or d != 1)
+
+    # --- stack ---
+    elif func_name == "stack":
+        # detect axis as last positional if candidate
+        elems = args[0]
+        if axis is None and len(args) > 1:
+            axis = args[1]
+        if axis is None:
+            axis = 0
+        return elems[0][:axis] + (len(elems),) + elems[0][axis:]
+
+    # --- tensordot ---
+    elif func_name == "tensordot":
+        if axes is None and len(args) > 2:
+            axes = args[2]
+        if axis is None:
+            axes = 2
+        if b is None:
+            return None
+        if isinstance(axes, int):
+            a_rest = a[:-axes]
+            b_rest = b[axes:]
+        else:
+            a_axes, b_axes = axes
+            a_rest = tuple(d for i, d in enumerate(a) if i not in a_axes)
+            b_rest = tuple(d for i, d in enumerate(b) if i not in b_axes)
+        return a_rest + b_rest
+
+    # --- transpose ---
+    elif func_name == ("transpose", "T", "mT"):
+        return a[:-2] + (a[-1], a[-2])
+
+    # --- vecdot ---
+    elif func_name == "vecdot":
+        if axis is None and len(args) > 2:
+            axis = args[2]
+        if axis is None:
+            axis = -1
+        if b is None:
+            return None
+        a_axis = axis + len(a)
+        b_axis = axis + len(b)
+        a_rem = tuple(d for i, d in enumerate(a) if i != a_axis)
+        b_rem = tuple(d for i, d in enumerate(b) if i != b_axis)
+        return broadcast_shapes(a_rem, b_rem)
+    else:
+        return None
+
+
 def reduce_shape(shape, axis, keepdims):
     """Reduce shape along given axis or axes (collapse dimensions)."""
     if shape is None:
@@ -133,8 +282,18 @@ def visit_Call(self, node):  # noqa : C901
             else:
                 kwargs[kw.arg] = self._lookup_value(kw.value)
 
+        # ------- handle linear algebra ---------------
+        target = None
         if func_name in lin_alg_funcs:
-            return None  # need to implement shape handling for these funcs
+            target = func_name
+        if attr_name in lin_alg_funcs:
+            target = attr_name
+        if target is not None:
+            args = [self.visit(arg) for arg in node.args]
+            # If it's a method call, prepend the object shape
+            if obj_shape is not None and attr_name == target:
+                args.insert(0, obj_shape)
+            return linalg_shape(target, args, kwargs)
 
         # ------- handle constructors ---------------
         if func_name in constructors or attr_name == "reshape":
@@ -241,7 +400,10 @@ def visit_Compare(self, node):
         return elementwise(*shapes)
 
     def visit_Constant(self, node):
-        return ()
+        return () if not hasattr(node.value, "shape") else node.value.shape
+
+    def visit_Tuple(self, node):
+        return tuple(self.visit(arg) for arg in node.elts)
 
     def visit_BinOp(self, node):
         left = self.visit(node.left)

tests/ndarray/test_lazyexpr.py

@@ -1129,10 +1129,10 @@ def test_rebasing(array_fixture):
     assert expr.expression == "(o0 + o1 - o2 * o3)"
 
     expr = blosc2.lazyexpr("a1")
-    assert expr.expression == "o0"
+    assert expr.expression == "(o0)"
 
     expr = blosc2.lazyexpr("a1[:10]")
-    assert expr.expression == "o0.slice((slice(None, 10, None),))"
+    assert expr.expression == "(o0.slice((slice(None, 10, None),)))"
 
 
 # Test get_chunk method
@@ -1595,7 +1595,96 @@ def test_not_numexpr():
     np.testing.assert_array_almost_equal(d_blosc2, np.logaddexp(npa, npb) + npa)
     # TODO: Implement __add__ etc. for LazyUDF so this line works
     # d_blosc2 = blosc2.evaluate(f"logaddexp(a, b) + clip(a, 6, 12)")
-    arr = blosc2.lazyexpr("matmul(a,b) + a ")
+    arr = blosc2.lazyexpr("matmul(a, b)")
     assert isinstance(arr, blosc2.LazyExpr)
-    assert arr.shape is None  # can't calculate shape for linalg funcs yet
-    np.testing.assert_array_almost_equal(arr[()], np.matmul(npa, npb) + a)
+    np.testing.assert_array_almost_equal(arr[()], np.matmul(npa, npb))
+
+
+def test_lazylinalg():
+    """
+    Test the shape parser for linear algebra funcs
+    """
+    # --- define base shapes ---
+    shapes = {
+        "A": (3, 4),
+        "B": (4, 5),
+        "C": (2, 3, 4),
+        "D": (1, 5, 1),
+        "x": (10,),
+        "y": (10,),
+    }
+    s = shapes["x"]
+    x = blosc2.linspace(0, np.prod(s), shape=s)
+    s = shapes["y"]
+    y = blosc2.linspace(0, np.prod(s), shape=s)
+    s = shapes["A"]
+    A = blosc2.linspace(0, np.prod(s), shape=s)
+    s = shapes["B"]
+    B = blosc2.linspace(0, np.prod(s), shape=s)
+    s = shapes["C"]
+    C = blosc2.linspace(0, np.prod(s), shape=s)
+    s = shapes["D"]
+    D = blosc2.linspace(0, np.prod(s), shape=s)
+
+    npx = x[()]
+    npy = y[()]
+    npA = A[()]
+
+    # --- concat ---
+    out = blosc2.lazyexpr("concat((x, y), axis=0)")
+    assert out.shape == np.concat((npx, npy), axis=0).shape
+
+    # --- diagonal ---
+    out = blosc2.lazyexpr("diagonal(A)")
+    assert out.shape == np.diagonal(npA).shape
+
+    # --- expand_dims ---
+    out = blosc2.lazyexpr("expand_dims(x, axis=0)")
+    assert out.shape == (1,) + shapes["x"]
+
+    # --- matmul ---
+    out = blosc2.lazyexpr("matmul(A, B)")
+    assert out.shape == (shapes["A"][0], shapes["B"][1])
+
+    # --- matrix_transpose ---
+    out = blosc2.lazyexpr("matrix_transpose(A)")
+    assert out.shape == (shapes["A"][1], shapes["A"][0])
+
+    # --- outer ---
+    out = blosc2.lazyexpr("outer(x, y)")
+    assert out.shape == shapes["x"] + shapes["y"]
+
+    # --- permute_dims ---
+    out = blosc2.lazyexpr("permute_dims(C, axes=(2,0,1))")
+    assert out.shape == (shapes["C"][2], shapes["C"][0], shapes["C"][1])
+
+    # --- squeeze ---
+    out = blosc2.lazyexpr("squeeze(D)")
+    assert out.shape == (5,)
+    out = blosc2.lazyexpr("D.squeeze()")
+    assert out.shape == (5,)
+
+    # --- stack ---
+    out = blosc2.lazyexpr("stack((x, y), axis=0)")
+    assert out.shape == (2,) + shapes["x"]
+
+    # --- tensordot ---
+    out = blosc2.lazyexpr("tensordot(A, B, axes=1)")
+    assert out.shape[0] == shapes["A"][0]
+    assert out.shape[-1] == shapes["B"][-1]
+
+    # --- vecdot ---
+    out = blosc2.lazyexpr("vecdot(x, y)")
+    assert out.shape == np.vecdot(x[()], y[()]).shape
+
+    # batched matmul
+    shapes = {
+        "A": (1, 3, 4),
+        "B": (3, 4, 5),
+    }
+    s = shapes["A"]
+    A = blosc2.linspace(0, np.prod(s), shape=s)
+    s = shapes["B"]
+    B = blosc2.linspace(0, np.prod(s), shape=s)
+    out = blosc2.lazyexpr("matmul(A, B)")
+    assert out.shape == (3, 3, 5)