Add Francesc's suggestions, improve tests

Author: lshaw8317

src/blosc2/lazyexpr.py

@@ -53,7 +53,7 @@
 
 from .shape_utils import constructors, infer_shape, lin_alg_funcs, reducers
 
-lin_alg_funcs += ("clip", "logaddexp")
+not_numexpr_funcs = lin_alg_funcs + ("clip", "logaddexp")
 
 if not blosc2.IS_WASM:
     import numexpr
@@ -645,8 +645,11 @@ def validate_expr(expr: str) -> None:
     skip_quotes = re.sub(r"(\'[^\']*\')", "", no_whitespace)
 
     # Check for forbidden patterns
-    if _blacklist_re.search(skip_quotes) is not None:
-        raise ValueError(f"'{expr}' is not a valid expression.")
+    forbiddens = _blacklist_re.search(skip_quotes)
+    if forbiddens is not None:
+        i = forbiddens.span()[0]
+        if expr[i : i + 2] != ".T" and expr[i : i + 3] != ".mT":  # allow tranpose methods
+            raise ValueError(f"'{expr}' is not a valid expression.")
 
     # Check for invalid characters not covered by the tokenizer
     invalid_chars = re.compile(r"[^\w\s+\-*/%()[].,=<>!&|~^]")
@@ -706,7 +709,7 @@ def visit_Call(self, node):
                 try:
                     shape = infer_shape(full_expr, shapes)
                 except Exception as e:
-                    print(f"⚠️ Shape inference failed for {full_expr}: {e}")
+                    print(f"Shape inference failed for {full_expr}: {e}")
                     shape = ()
 
                 # Determine dtype
@@ -2920,7 +2923,7 @@ def find_args(expr):
         return value, expression[idx:idx2]
 
     def _compute_expr(self, item, kwargs):
-        if any(method in self.expression for method in reducers + lin_alg_funcs):
+        if any(method in self.expression for method in reducers + not_numexpr_funcs):
             # We have reductions in the expression (probably coming from a string lazyexpr)
             # Also includes slice
             _globals = get_expr_globals(self.expression)
@@ -3441,8 +3444,8 @@ def _numpy_eval_expr(expression, operands, prefer_blosc=False):
     if "contains" in expression:
         _out = ne_evaluate(expression, local_dict=ops)
     else:
-        # Create a globals dict with the functions of blosc2 preferentially
-        # (and numpy if can't find blosc2)
+        # Create a globals dict with blosc2 version of functions preferentially
+        # (default to numpy func if not implemented in blosc2)
         if prefer_blosc:
             _globals = get_expr_globals(expression)
             _globals |= dtype_symbols

src/blosc2/shape_utils.py

@@ -104,9 +104,9 @@ def linalg_shape(func_name, args, kwargs):  # noqa: C901
         batch = broadcast_shapes(a[:-2], b[:-2])
         shape = batch
         if not x1_is_vector:
-            shape += a[-2]
+            shape += (a[-2],)
         if not x2_is_vector:
-            shape += b[-1]
+            shape += (b[-1],)
         return shape
 
     # --- matrix_transpose ---
@@ -154,7 +154,7 @@ def linalg_shape(func_name, args, kwargs):  # noqa: C901
     elif func_name == "tensordot":
         if axes is None and len(args) > 2:
             axes = args[2]
-        if axis is None:
+        if axes is None:
             axes = 2
         if b is None:
             return None
@@ -168,7 +168,7 @@ def linalg_shape(func_name, args, kwargs):  # noqa: C901
         return a_rest + b_rest
 
     # --- transpose ---
-    elif func_name == ("transpose", "T", "mT"):
+    elif func_name in ("transpose", "T", "mT"):
         return a[:-2] + (a[-1], a[-2])
 
     # --- vecdot ---
@@ -261,9 +261,22 @@ def visit_Name(self, node):
         else:  # passed a scalar value
             return ()
 
+    def visit_Attribute(self, node):
+        obj_shape = self.visit(node.value)
+        attr = node.attr
+        if attr == "reshape":
+            if node.args:
+                shape_arg = node.args[-1]
+                if isinstance(shape_arg, ast.Tuple):
+                    return tuple(self._lookup_value(e) for e in shape_arg.elts)
+            return ()
+        elif attr in ("T", "mT"):
+            return linalg_shape(attr, (obj_shape,), {})
+        return None
+
     def visit_Call(self, node):  # noqa : C901
         func_name = getattr(node.func, "id", None)
-        attr_name = getattr(node.func, "attr", None)
+        attr_name = getattr(node.func, "attr", None)  # handle methods called on funcs
 
         # --- Recursive method-chain support ---
         obj_shape = None

tests/ndarray/test_lazyexpr.py

@@ -1629,62 +1629,102 @@ def test_lazylinalg():
     npx = x[()]
     npy = y[()]
     npA = A[()]
+    npB = B[()]
+    npC = C[()]
+    npD = D[()]
 
     # --- concat ---
     out = blosc2.lazyexpr("concat((x, y), axis=0)")
-    assert out.shape == np.concat((npx, npy), axis=0).shape
+    npres = np.concatenate((npx, npy), axis=0)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- diagonal ---
     out = blosc2.lazyexpr("diagonal(A)")
-    assert out.shape == np.diagonal(npA).shape
+    npres = np.diagonal(npA)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- expand_dims ---
     out = blosc2.lazyexpr("expand_dims(x, axis=0)")
-    assert out.shape == (1,) + shapes["x"]
+    npres = np.expand_dims(npx, axis=0)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- matmul ---
     out = blosc2.lazyexpr("matmul(A, B)")
-    assert out.shape == (shapes["A"][0], shapes["B"][1])
+    npres = np.matmul(npA, npB)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- matrix_transpose ---
     out = blosc2.lazyexpr("matrix_transpose(A)")
-    assert out.shape == (shapes["A"][1], shapes["A"][0])
+    npres = np.matrix_transpose(npA)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
+    out = blosc2.lazyexpr("C.mT")
+    npres = C.mT
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
+    out = blosc2.lazyexpr("A.T")
+    npres = npA.T
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- outer ---
     out = blosc2.lazyexpr("outer(x, y)")
-    assert out.shape == shapes["x"] + shapes["y"]
+    npres = np.outer(npx, npy)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- permute_dims ---
     out = blosc2.lazyexpr("permute_dims(C, axes=(2,0,1))")
-    assert out.shape == (shapes["C"][2], shapes["C"][0], shapes["C"][1])
+    npres = np.transpose(npC, axes=(2, 0, 1))
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- squeeze ---
     out = blosc2.lazyexpr("squeeze(D)")
-    assert out.shape == (5,)
+    npres = np.squeeze(npD)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
+
     out = blosc2.lazyexpr("D.squeeze()")
-    assert out.shape == (5,)
+    npres = np.squeeze(npD)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- stack ---
     out = blosc2.lazyexpr("stack((x, y), axis=0)")
-    assert out.shape == (2,) + shapes["x"]
+    npres = np.stack((npx, npy), axis=0)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- tensordot ---
     out = blosc2.lazyexpr("tensordot(A, B, axes=1)")
-    assert out.shape[0] == shapes["A"][0]
-    assert out.shape[-1] == shapes["B"][-1]
+    npres = np.tensordot(npA, npB, axes=1)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
     # --- vecdot ---
     out = blosc2.lazyexpr("vecdot(x, y)")
-    assert out.shape == np.vecdot(x[()], y[()]).shape
+    npres = np.vecdot(npx, npy)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)
 
-    # batched matmul
+    # --- batched matmul ---
     shapes = {
         "A": (1, 3, 4),
         "B": (3, 4, 5),
     }
     s = shapes["A"]
     A = blosc2.linspace(0, np.prod(s), shape=s)
+    npA = A[()]  # actual numpy array
     s = shapes["B"]
     B = blosc2.linspace(0, np.prod(s), shape=s)
+    npB = B[()]  # actual numpy array
+
     out = blosc2.lazyexpr("matmul(A, B)")
-    assert out.shape == (3, 3, 5)
+    npres = np.matmul(npA, npB)
+    assert out.shape == npres.shape
+    np.testing.assert_array_almost_equal(out[()], npres)