Add UnaryOp handling to infer_shape

Author: lshaw8317

src/blosc2/lazyexpr.py

@@ -47,6 +47,7 @@
     process_key,
 )
 
+from .proxy import _convert_dtype
 from .shape_utils import constructors, elementwise_funcs, infer_shape, linalg_attrs, linalg_funcs, reducers
 
 if not blosc2.IS_WASM:
@@ -2213,7 +2214,9 @@ def result_type(
     # Follow NumPy rules for scalar-array operations
     # Create small arrays with the same dtypes and let NumPy's type promotion determine the result type
     arrs = [
-        value if (np.isscalar(value) or not hasattr(value, "dtype")) else np.array([0], dtype=value.dtype)
+        value
+        if (np.isscalar(value) or not hasattr(value, "dtype"))
+        else np.array([0], dtype=_convert_dtype(value.dtype))
         for value in arrays_and_dtypes
     ]
     return np.result_type(*arrs)

src/blosc2/proxy.py

@@ -579,6 +579,8 @@ def _convert_dtype(dt: str | DTypeLike):
     """
     Attempts to convert to blosc2.dtype (i.e. numpy dtype)
     """
+    if hasattr(dt, "as_numpy_dtype"):
+        dt = dt.as_numpy_dtype
     try:
         return np.dtype(dt)
     except TypeError:  # likely passed e.g. a torch.float64
@@ -616,7 +618,7 @@ def __init__(self, src, chunks: tuple | None = None, blocks: tuple | None = None
             raise TypeError("The source must have a __getitem__ method")
         self._src = src
         self._dtype = _convert_dtype(src.dtype)
-        self._shape = src.shape
+        self._shape = src.shape if isinstance(src.shape, tuple) else tuple(src.shape)
         # Compute reasonable values for chunks and blocks
         cparams = blosc2.CParams(clevel=0)
 

src/blosc2/shape_utils.py

@@ -520,6 +520,9 @@ def visit_BinOp(self, node):
         right = self.visit(node.right)
         return elementwise(left, right)
 
+    def visit_UnaryOp(self, node):
+        return self.visit(node.operand)
+
     def _eval_slice(self, node):
         if isinstance(node, ast.Slice):
             return slice(

tests/ndarray/test_lazyexpr.py

@@ -1195,7 +1195,7 @@ def test_fill_disk_operands(chunks, blocks, disk, fill_value):
         b = blosc2.open("b.b2nd")
         c = blosc2.open("c.b2nd")
 
-    expr = ((a**3 + blosc2.sin(c * 2)) < b) & (c > 0)
+    expr = ((a**3 + blosc2.sin(c * 2)) < b) & ~(c > 0)
 
     out = expr.compute()
     assert out.shape == (N, N)

tests/ndarray/test_proxy.py

@@ -6,8 +6,15 @@
 # LICENSE file in the root directory of this source tree)
 #######################################################################
 
+import dask.array as da
+import jax.numpy as jnp
 import numpy as np
 import pytest
+import tensorflow as tf
+import torch
+
+# TODO: import cupy as cp
+import zarr
 
 import blosc2
 from blosc2.ndarray import get_chunks_idx
@@ -172,3 +179,52 @@ def get_chunk(self, nchunk):
     proxy = blosc2.Proxy(source)
     result = proxy[...]
     np.testing.assert_array_equal(result, data)
+
+
+@pytest.mark.parametrize(
+    "xp",
+    [torch, tf, np, jnp, da, zarr],
+)
+@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)