Ensure casting to SimpleProxy for non-string expressions

Author: Luke Shaw

src/blosc2/lazyexpr.py

@@ -2258,6 +2258,12 @@ def __init__(self, new_op):  # noqa: C901
             return
         value1, op, value2 = new_op
         dtype_ = check_dtype(op, value1, value2)  # perform some checks
+        # Check that operands are proper Operands, LazyArray or scalars; if not, convert to NDArray objects
+        value1 = (
+            blosc2.SimpleProxy(value1)
+            if not (isinstance(value1, (blosc2.Operand, np.ndarray)) or np.isscalar(value1))
+            else value1
+        )
         if value2 is None:
             if isinstance(value1, LazyExpr):
                 self.expression = value1.expression if op is None else f"{op}({value1.expression})"
@@ -2266,7 +2272,12 @@ def __init__(self, new_op):  # noqa: C901
                 self.operands = {"o0": value1}
                 self.expression = "o0" if op is None else f"{op}(o0)"
             return
-        elif isinstance(value1, LazyExpr) or isinstance(value2, LazyExpr):
+        value2 = (
+            blosc2.SimpleProxy(value2)
+            if not (isinstance(value2, (blosc2.Operand, np.ndarray)) or np.isscalar(value2))
+            else value2
+        )
+        if isinstance(value1, LazyExpr) or isinstance(value2, LazyExpr):
             if isinstance(value1, LazyExpr):
                 newexpr = value1.update_expr(new_op)
             else:
@@ -2739,7 +2750,7 @@ def find_args(expr):
 
     def _compute_expr(self, item, kwargs):  # noqa : C901
         # ne_evaluate will need safe_blosc2_globals for some functions (e.g. clip, logaddexp)
-        # that are implemenetd in python-blosc2 not in numexpr
+        # that are implemented in python-blosc2 not in numexpr
         global safe_blosc2_globals
         if len(safe_blosc2_globals) == 0:
             # First eval call, fill blosc2_safe_globals for ne_evaluate
@@ -3011,7 +3022,7 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
             _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():
-                if not (isinstance(val, (blosc2.Operand, blosc2.LazyArray, np.ndarray)) or np.isscalar(val)):
+                if not (isinstance(val, (blosc2.Operand, np.ndarray)) or np.isscalar(val)):
                     _operands[op] = blosc2.SimpleProxy(val)
             # 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()}

src/blosc2/linalg.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 import blosc2
-from blosc2.ndarray import get_intersecting_chunks, npvecdot, slice_to_chunktuple
+from blosc2.ndarray import get_intersecting_chunks, nptranspose, npvecdot, slice_to_chunktuple
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -79,9 +79,8 @@ def matmul(x1: blosc2.Array, x2: blosc2.NDArray, **kwargs: Any) -> blosc2.NDArra
     if np.isscalar(x1) or np.isscalar(x2):
         raise ValueError("Arguments can't be scalars.")
 
-    # Added this to pass array-api tests (which use internal getitem to check results)
-    x1 = blosc2.asarray(x1)
-    x2 = blosc2.asarray(x2)
+    # Makes a SimpleProxy if inputs are not blosc2 arrays
+    x1, x2 = blosc2.asarray(x1), blosc2.asarray(x2)
 
     # Validate matrix multiplication compatibility
     if x1.shape[builtins.max(-1, -len(x2.shape))] != x2.shape[builtins.max(-2, -len(x2.shape))]:
@@ -183,9 +182,6 @@ def tensordot(
     """
     fast_path = kwargs.pop("fast_path", None)  # for testing purposes
     # TODO: add fast path for when don't need to change chunkshapes
-    # 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.tensordot(x1, x2, axes=axes)
 
     x1, x2 = blosc2.asarray(x1), blosc2.asarray(x2)
 
@@ -261,24 +257,8 @@ def tensordot(
         a_selection = tuple(next(rchunk_iter) if a else slice(None, None, 1) for a in a_keep)
         b_selection = tuple(next(rchunk_iter) if b else slice(None, None, 1) for b in b_keep)
         res_chunks = tuple(s.stop - s.start for s in res_chunk)
-
-        if fast_path:  # just load everything
-            bx1 = x1[a_selection]
-            bx2 = x2[b_selection]
-            newshape_a = (
-                math.prod([bx1.shape[i] for i in a_keep_axes]),
-                math.prod([bx1.shape[a] for a in a_axes]),
-            )
-            newshape_b = (
-                math.prod([bx2.shape[b] for b in b_axes]),
-                math.prod([bx2.shape[i] for i in b_keep_axes]),
-            )
-            at = bx1.transpose(newaxes_a).reshape(newshape_a)
-            bt = bx2.transpose(newaxes_b).reshape(newshape_b)
-            res = np.dot(at, bt)
-            result[res_chunk] += res.reshape(res_chunks)
-        else:  # operands too big, have to go chunk-by-chunk
-            for ochunk in product(*op_chunks):
+        for ochunk in product(*op_chunks):
+            if not fast_path:  # operands too big, have to go chunk-by-chunk
                 op_chunk = tuple(
                     slice(rc * rcs, builtins.min((rc + 1) * rcs, x1s), 1)
                     for rc, rcs, x1s in zip(ochunk, a_chunks_red, a_shape_red, strict=True)
@@ -293,21 +273,23 @@ def tensordot(
                     op_chunk[next(order_iter)] if not b else bs_
                     for bs_, b in zip(b_selection, b_keep, strict=True)
                 )
-                bx1 = x1[a_selection]
-                bx2 = x2[b_selection]
-                # adapted from numpy tensordot
-                newshape_a = (
-                    math.prod([bx1.shape[i] for i in a_keep_axes]),
-                    math.prod([bx1.shape[a] for a in a_axes]),
-                )
-                newshape_b = (
-                    math.prod([bx2.shape[b] for b in b_axes]),
-                    math.prod([bx2.shape[i] for i in b_keep_axes]),
-                )
-                at = bx1.transpose(newaxes_a).reshape(newshape_a)
-                bt = bx2.transpose(newaxes_b).reshape(newshape_b)
-                res = np.dot(at, bt)
-                result[res_chunk] += res.reshape(res_chunks)
+            bx1 = x1[a_selection]
+            bx2 = x2[b_selection]
+            # adapted from numpy tensordot
+            newshape_a = (
+                math.prod([bx1.shape[i] for i in a_keep_axes]),
+                math.prod([bx1.shape[a] for a in a_axes]),
+            )
+            newshape_b = (
+                math.prod([bx2.shape[b] for b in b_axes]),
+                math.prod([bx2.shape[i] for i in b_keep_axes]),
+            )
+            at = nptranspose(bx1, newaxes_a).reshape(newshape_a)
+            bt = nptranspose(bx2, newaxes_b).reshape(newshape_b)
+            res = np.dot(at, bt)
+            result[res_chunk] += res.reshape(res_chunks)
+            if fast_path:  # already done everything
+                break
     return result
 
 
@@ -396,19 +378,17 @@ def vecdot(x1: blosc2.NDArray, x2: blosc2.NDArray, axis: int = -1, **kwargs) ->
         )
         b_selection = tuple(next(rchunk_iter) if b else slice(None, None, 1) for b in b_keep)
 
-        if fast_path:  # just load everything, also handles case of 0 in shapes
-            bx1 = x1[a_selection]
-            bx2 = x2[b_selection]
-            result[res_chunk] += npvecdot(bx1, bx2, axis=axis)  # handles conjugation of bx1
-        else:  # operands too big, have to go chunk-by-chunk
-            for ochunk in range(0, a_shape_red, a_chunks_red):
+        for ochunk in range(0, a_shape_red, a_chunks_red):
+            if not fast_path:  # operands too big, go chunk-by-chunk
                 op_chunk = (slice(ochunk, builtins.min(ochunk + a_chunks_red, x1.shape[a_axes]), 1),)
                 a_selection = a_selection[:a_axes] + op_chunk + a_selection[a_axes + 1 :]
                 b_selection = b_selection[:b_axes] + op_chunk + b_selection[b_axes + 1 :]
-                bx1 = x1[a_selection]
-                bx2 = x2[b_selection]
-                res = npvecdot(bx1, bx2, axis=axis)  # handles conjugation of bx1
-                result[res_chunk] += res
+            bx1 = x1[a_selection]
+            bx2 = x2[b_selection]
+            res = npvecdot(bx1, bx2, axis=axis)  # handles conjugation of bx1
+            result[res_chunk] += res
+            if fast_path:  # already done everything
+                break
     return result
 
 
@@ -517,7 +497,7 @@ def permute_dims(
         src_slice = tuple(slice(start, stop) for start, stop in start_stop)
         dst_slice = tuple(slice(start_stop[ax][0], start_stop[ax][1]) for ax in axes)
 
-        transposed = np.transpose(arr[src_slice], axes=axes)
+        transposed = nptranspose(arr[src_slice], axes=axes)
         result[dst_slice] = np.ascontiguousarray(transposed)
 
     return result
@@ -648,6 +628,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.asarray(x1), blosc2.asarray(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

@@ -42,10 +42,12 @@
     nprshift = np.bitwise_right_shift
     npbinvert = np.bitwise_invert
     npvecdot = np.vecdot
+    nptranspose = np.permute_dims
 else:  # not array-api compliant
     nplshift = np.left_shift
     nprshift = np.right_shift
     npbinvert = np.bitwise_not
+    nptranspose = np.transpose
 
     def npvecdot(a, b, axis=-1):
         return np.einsum("...i,...i->...", np.moveaxis(np.conj(a), axis, -1), np.moveaxis(b, axis, -1))
@@ -2969,7 +2971,8 @@ def chunkwise_clip(inputs, output, offset):
         x, min, max = inputs
         output[:] = np.clip(x, min, max)
 
-    return blosc2.lazyudf(chunkwise_clip, (x, min, max), dtype=x.dtype, shape=x.shape, **kwargs)
+    dtype = blosc2.result_type(x)
+    return blosc2.lazyudf(chunkwise_clip, (x, min, max), dtype=dtype, shape=x.shape, **kwargs)
 
 
 def logaddexp(x1: int | float | blosc2.Array, x2: int | float | blosc2.Array, **kwargs: Any) -> NDArray:
@@ -3001,7 +3004,7 @@ def chunkwise_logaddexp(inputs, output, offset):
         x1, x2 = inputs
         output[:] = np.logaddexp(x1, x2)
 
-    dtype = blosc2.result_type(x1.dtype, x2.dtype)
+    dtype = blosc2.result_type(x1, x2)
     if dtype == blosc2.bool_:
         raise TypeError("logaddexp doesn't accept boolean arguments.")
 
@@ -5653,7 +5656,8 @@ def asarray(array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs:
         raise ValueError("Only unsafe casting is supported at the moment.")
     if not hasattr(array, "shape"):
         array = np.asarray(array)  # defaults if dtype=None
-    dtype = kwargs.pop("dtype", array.dtype)  # check if dtype provided
+    dtype_ = blosc2.proxy._convert_dtype(array.dtype)
+    dtype = kwargs.pop("dtype", dtype_)  # check if dtype provided
     kwargs = _check_ndarray_kwargs(**kwargs)
     chunks = kwargs.pop("chunks", None)
     blocks = kwargs.pop("blocks", None)
@@ -5664,14 +5668,13 @@ def asarray(array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs:
     # Let's avoid this
     if blocks is None and hasattr(array, "blocks") and isinstance(array.blocks, (tuple, list)):
         blocks = array.blocks
-    chunks, blocks = compute_chunks_blocks(array.shape, chunks, blocks, array.dtype, **kwargs)
 
-    copy = True if copy is None and not isinstance(array, NDArray) else copy
     if copy:
+        chunks, blocks = compute_chunks_blocks(array.shape, chunks, blocks, dtype_, **kwargs)
         # Fast path for small arrays. This is not too expensive in terms of memory consumption.
         shape = array.shape
         small_size = 2**24  # 16 MB
-        array_nbytes = math.prod(shape) * array.dtype.itemsize
+        array_nbytes = math.prod(shape) * dtype_.itemsize
         if array_nbytes < small_size:
             if not isinstance(array, np.ndarray) and hasattr(array, "chunks"):
                 # A getitem operation should be enough to get a numpy array
@@ -5682,7 +5685,7 @@ def asarray(array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs:
             return blosc2_ext.asarray(array, chunks, blocks, **kwargs)
 
         # Create the empty array
-        ndarr = empty(shape, array.dtype, chunks=chunks, blocks=blocks, **kwargs)
+        ndarr = empty(shape, dtype_, chunks=chunks, blocks=blocks, **kwargs)
         behaved = are_partitions_behaved(shape, chunks, blocks)
 
         # Get the coordinates of the chunks
@@ -5705,7 +5708,7 @@ def asarray(array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs:
                 ndarr[slice_] = array_slice
     else:
         if not isinstance(array, NDArray):
-            raise ValueError("Must always do a copy for asarray unless NDArray provided.")
+            return blosc2.SimpleProxy(array, chunks, blocks)
         # TODO: make a direct view possible
         return array
 

src/blosc2/proxy.py

@@ -649,6 +649,11 @@ def dtype(self):
         """The data type of the source array."""
         return self._dtype
 
+    @property
+    def ndim(self):
+        """The number of dimensions of the source array."""
+        return len(self.shape)
+
     def __getitem__(self, item: slice | list[slice]) -> np.ndarray:
         """
         Get a slice as a numpy.ndarray (via this proxy).