Handle foreign inputs to logaddexp, clip

Author: lshaw8317

ADD_LAZYFUNCS.md

@@ -5,6 +5,8 @@ Once you have written a (public API) function in Blosc2, it is important to:
 * Add it to the list of functions in ``__all__`` in the ``__init__.py`` file
 * If it is present in numpy, add it to the relevant dictionary (``local_ufunc_map``, ``ufunc_map`` ``ufunc_map_1param``) in ``ndarray.py``
 
+If your function is implemented at the Blosc2 level (and not via either the `LazyUDF` or `LazyExpr` classes), you will need to add some conversion of the inputs to SimpleProxy instances (see e.g. ``matmul`` for an example).
+
 Finally, you also need to deal with it correctly within ``shape_utils.py``.
 
 If the function does not change the shape of the output, simply add it to ``elementwise_funcs`` and you're done.

src/blosc2/__init__.py

@@ -445,7 +445,7 @@ def _raise(exc):
     result_type,
     can_cast,
 )
-from .proxy import Proxy, ProxySource, ProxyNDSource, ProxyNDField, SimpleProxy, jit
+from .proxy import Proxy, ProxySource, ProxyNDSource, ProxyNDField, SimpleProxy, jit, as_simpleproxy
 
 from .schunk import SChunk, open
 from . import linalg
@@ -648,6 +648,7 @@ def _raise(exc):
     "asarray",
     "asin",
     "asinh",
+    "as_simpleproxy",
     "astype",
     "atan",
     "atan2",

src/blosc2/lazyexpr.py

@@ -434,13 +434,13 @@ def convert_inputs(inputs):
         return []
     inputs_ = []
     for obj in inputs:
-        if not isinstance(obj, blosc2.Array) and not np.isscalar(obj):
+        if not isinstance(obj, (np.ndarray, blosc2.Operand)) and not np.isscalar(obj):
             try:
-                obj = np.asarray(obj)
+                obj = blosc2.SimpleProxy(obj)
             except Exception:
                 print(
                     "Inputs not being np.ndarray, Array or Python scalar objects"
-                    " should be convertible to np.ndarray."
+                    " should be convertible to SimpleProxy."
                 )
                 raise
         inputs_.append(obj)
@@ -1077,9 +1077,9 @@ def fill_chunk_operands(  # noqa: C901
         if nchunk == 0:
             # Initialize the iterator for reading the chunks
             # Take any operand (all should have the same shape and chunks)
-            arr = next(iter(operands.values()))
+            key, arr = next(iter(operands.items()))
             chunks_idx, _ = get_chunks_idx(arr.shape, arr.chunks)
-            info = (reduc, aligned, low_mem, chunks_idx)
+            info = (reduc, aligned[key], low_mem, chunks_idx)
             iter_chunks = read_nchunk(list(operands.values()), info)
         # Run the asynchronous file reading function from a synchronous context
         chunks = next(iter_chunks)
@@ -1095,7 +1095,7 @@ def fill_chunk_operands(  # noqa: C901
                 # The chunk is a special zero chunk, so we can treat it as a scalar
                 chunk_operands[key] = np.zeros((), dtype=value.dtype)
                 continue
-            if aligned:
+            if aligned[key]:
                 buff = blosc2.decompress2(chunks[i])
                 bsize = value.dtype.itemsize * math.prod(chunks_)
                 chunk_operands[key] = np.frombuffer(buff[:bsize], dtype=value.dtype).reshape(chunks_)
@@ -1115,10 +1115,6 @@ def fill_chunk_operands(  # noqa: C901
             chunk_operands[key] = value[()]
             continue
 
-        if isinstance(value, np.ndarray | blosc2.C2Array):
-            chunk_operands[key] = value[slice_]
-            continue
-
         if not full_chunk or not isinstance(value, blosc2.NDArray):
             # The chunk is not a full one, or has padding, or is not a blosc2.NDArray,
             # so we need to go the slow path
@@ -1143,7 +1139,7 @@ def fill_chunk_operands(  # noqa: C901
             value.get_slice_numpy(chunk_operands[key], (starts, stops))
             continue
 
-        if aligned:
+        if aligned[key]:
             # Decompress the whole chunk and store it
             buff = value.schunk.decompress_chunk(nchunk)
             bsize = value.dtype.itemsize * math.prod(chunks_)
@@ -1203,7 +1199,10 @@ def fast_eval(  # noqa: C901
     if blocks is None:
         blocks = basearr.blocks
     # Check whether the partitions are aligned and behaved
-    aligned = blosc2.are_partitions_aligned(shape, chunks, blocks)
+    aligned = {
+        k: False if not hasattr(k, "chunks") else blosc2.are_partitions_aligned(k.shape, k.chunks, k.blocks)
+        for k in operands
+    }
     behaved = blosc2.are_partitions_behaved(shape, chunks, blocks)
 
     # Check that all operands are NDArray for fast path
@@ -1227,7 +1226,7 @@ def fast_eval(  # noqa: C901
         offset = tuple(s.start for s in cslice)  # offset for the udf
         chunks_ = tuple(s.stop - s.start for s in cslice)
 
-        full_chunk = chunks_ == chunks
+        full_chunk = chunks_ == chunks  # slice is same as chunk
         fill_chunk_operands(
             operands, cslice, chunks_, full_chunk, aligned, nchunk, iter_disk, chunk_operands
         )
@@ -1811,7 +1810,7 @@ def reduce_slices(  # noqa: C901
         same_chunks = all(operand.chunks == o.chunks for o in operands.values() if hasattr(o, "chunks"))
         same_blocks = all(operand.blocks == o.blocks for o in operands.values() if hasattr(o, "blocks"))
         fast_path = same_shape and same_chunks and same_blocks and (0 not in operand.chunks)
-        aligned, iter_disk = False, False
+        aligned, iter_disk = dict.fromkeys(operands.keys(), False), False
         if fast_path:
             chunks = operand.chunks
             # Check that all operands are NDArray for fast path

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.asarray(x1), blosc2.asarray(x2)
+    x1, x2 = blosc2.as_simpleproxy(x1), blosc2.as_simpleproxy(x2)
 
     # Validate matrix multiplication compatibility
     if x1.shape[builtins.max(-1, -len(x2.shape))] != x2.shape[builtins.max(-2, -len(x2.shape))]:
@@ -183,7 +183,8 @@ def tensordot(
     fast_path = kwargs.pop("fast_path", None)  # for testing purposes
     # TODO: add fast path for when don't need to change chunkshapes
 
-    x1, x2 = blosc2.asarray(x1), blosc2.asarray(x2)
+    # Makes a SimpleProxy if inputs are not blosc2 arrays
+    x1, x2 = blosc2.as_simpleproxy(x1), blosc2.as_simpleproxy(x2)
 
     if isinstance(axes, tuple):
         a_axes, b_axes = axes
@@ -324,7 +325,8 @@ def vecdot(x1: blosc2.NDArray, x2: blosc2.NDArray, axis: int = -1, **kwargs) ->
     if isinstance(x1, np.ndarray) and isinstance(x2, np.ndarray):
         return npvecdot(x1, x2, axis=axis)
 
-    x1, x2 = blosc2.asarray(x1), blosc2.asarray(x2)
+    # Makes a SimpleProxy if inputs are not blosc2 arrays
+    x1, x2 = blosc2.as_simpleproxy(x1), blosc2.as_simpleproxy(x2)
 
     N = builtins.min(x1.ndim, x2.ndim)
     if axis < -N or axis > -1:
@@ -466,8 +468,9 @@ def permute_dims(
     """
     if np.isscalar(arr) or arr.ndim < 2:
         return arr
-    if isinstance(arr, np.ndarray):  # for array-api test compliance (does getitem for comparison)
-        return np.permute_dims(arr, axes)
+
+    # Makes a SimpleProxy if input is not blosc2 array
+    arr = blosc2.as_simpleproxy(arr)
 
     ndim = arr.ndim
 
@@ -535,7 +538,8 @@ def transpose(x, **kwargs: Any) -> blosc2.NDArray:
     # If arguments are dimension < 2, they are returned
     if np.isscalar(x) or x.ndim < 2:
         return x
-
+    # Makes a SimpleProxy if input is not blosc2 array
+    x = blosc2.as_simpleproxy(x)
     # Validate arguments are dimension 2
     if x.ndim > 2:
         raise ValueError("Transposing arrays with dimension greater than 2 is not supported yet.")
@@ -559,6 +563,8 @@ def matrix_transpose(arr: blosc2.Array, **kwargs: Any) -> blosc2.NDArray:
         ``(..., N, M)``.
     """
     axes = None
+    # Makes a SimpleProxy if input is not blosc2 array
+    arr = blosc2.as_simpleproxy(arr)
     if not np.isscalar(arr) and arr.ndim > 2:
         axes = list(range(arr.ndim))
         axes[-2], axes[-1] = axes[-1], axes[-2]
@@ -592,6 +598,8 @@ def diagonal(x: blosc2.blosc2.NDArray, offset: int = 0) -> blosc2.blosc2.NDArray
 
     Reference: https://data-apis.org/array-api/latest/extensions/generated/array_api.linalg.diag.html#diag
     """
+    # Makes a SimpleProxy if input is not blosc2 array
+    x = blosc2.as_simpleproxy(x)
     n_rows, n_cols = x.shape[-2:]
     min_idx = builtins.min(n_rows, n_cols)
     if offset < 0:
@@ -628,7 +636,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)
+    x1, x2 = blosc2.as_simpleproxy(x1), blosc2.as_simpleproxy(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

@@ -5669,6 +5669,7 @@ def asarray(array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs:
     if blocks is None and hasattr(array, "blocks") and isinstance(array.blocks, (tuple, list)):
         blocks = array.blocks
 
+    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.
@@ -5708,7 +5709,7 @@ def asarray(array: Sequence | blosc2.Array, copy: bool | None = None, **kwargs:
                 ndarr[slice_] = array_slice
     else:
         if not isinstance(array, NDArray):
-            return blosc2.SimpleProxy(array, chunks, blocks)
+            raise ValueError("Must always do a copy for asarray unless NDArray provided.")
         # TODO: make a direct view possible
         return array
 

src/blosc2/proxy.py

@@ -670,6 +670,27 @@ def __getitem__(self, item: slice | list[slice]) -> np.ndarray:
         return np.asarray(self._src[item])  # avoids copy for PyTorch at least
 
 
+def as_simpleproxy(x: blosc2.Array) -> 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.
+
+    Returns
+    -------
+    out: blosc2.SimpleProxy | blosc2.Operand
+        Object with minimal interface for blosc2 LazyExpr computations.
+    """
+    if isinstance(x, blosc2.Operand):
+        return x
+    else:
+        return SimpleProxy(x)
+
+
 def jit(func=None, *, out=None, disable=False, **kwargs):  # noqa: C901
     """
     Prepare a function so that it can be used with the Blosc2 compute engine.

tests/ndarray/test_elementwise_funcs.py

@@ -188,7 +188,7 @@ def _test_unary_func_proxy(np_func, blosc_func, dtype, shape, xp):
         a_blosc[tuple(i // 4 for i in shape)] = 1 + 1j
         a_blosc[tuple(i // 2 for i in shape)] = xp.nan + xp.nan * 1j
     if dtype == blosc2.bool_ and np_func.__name__ == "arctanh":
-        a_blosc = xp.zeros(shape=shape, dtype=dtype_)
+        a_blosc = xp.zeros(shape, dtype=dtype_)
 
     arr = np.asarray(a_blosc)
     success = False
@@ -223,7 +223,7 @@ def _test_binary_func_impl(np_func, blosc_func, dtype, shape, chunkshape):  # no
         1, stop=np.prod(shape), num=np.prod(shape), chunks=chunkshape, shape=shape, dtype=dtype
     )
     if np_func.__name__ in ("right_shift", "left_shift"):
-        a_blosc2 = blosc2.asarray(2)
+        a_blosc2 = blosc2.asarray(2, copy=True)
     else:
         a_blosc2 = blosc2.linspace(
             start=np.prod(shape) * 2,
@@ -304,8 +304,8 @@ def test_unary_funcs(np_func, blosc_func, dtype, shape, chunkshape):
 @pytest.mark.parametrize("dtype", STR_DTYPES)
 @pytest.mark.parametrize("shape", [(10,), (20, 20)])
 @pytest.mark.parametrize("xp", [torch])
-def test_unfuncs_proxy(np_func, blosc_func, dtype, shape, chunkshape, xp):
-    _test_unary_func_proxy(np_func, blosc_func, dtype, shape, chunkshape, xp)
+def test_unfuncs_proxy(np_func, blosc_func, dtype, shape, xp):
+    _test_unary_func_proxy(np_func, blosc_func, dtype, shape, xp)
 
 
 @pytest.mark.heavy