diff --git a/ADD_LAZYFUNCS.md b/ADD_LAZYFUNCS.md index 4807d5dc..2c8175b0 100644 --- a/ADD_LAZYFUNCS.md +++ b/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. diff --git a/src/blosc2/__init__.py b/src/blosc2/__init__.py index 49d99ef5..ec2fc71b 100644 --- a/src/blosc2/__init__.py +++ b/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", diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 7891b506..1c6ba9fe 100644 --- a/src/blosc2/lazyexpr.py +++ b/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: @@ -433,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) @@ -885,7 +886,13 @@ def validate_inputs(inputs: dict, out=None, reduce=False) -> tuple: # noqa: C90 return shape, None, None, False # More checks specific of NDArray inputs - NDinputs = [input for input in inputs if hasattr(input, "chunks")] + # NDInputs are either non-SimpleProxy with chunks or are SimpleProxy with src having chunks + NDinputs = [ + input + for input in inputs + if (hasattr(input, "chunks") and not isinstance(input, blosc2.SimpleProxy)) + or (isinstance(input, blosc2.SimpleProxy) and hasattr(input.src, "chunks")) + ] if not NDinputs: # All inputs are NumPy arrays, so we cannot take the fast path if inputs and hasattr(inputs[0], "shape"): @@ -1076,9 +1083,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) @@ -1094,7 +1101,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_) @@ -1114,10 +1121,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 @@ -1142,7 +1145,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_) @@ -1202,7 +1205,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 @@ -1226,7 +1232,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 ) @@ -1810,7 +1816,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 @@ -2213,7 +2219,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) @@ -2255,15 +2263,29 @@ 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})" self.operands = value1.operands else: + if np.isscalar(value1): + value1 = ne_evaluate(f"{op}({value1})") + op = None 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: @@ -2274,7 +2296,8 @@ def __init__(self, new_op): # noqa: C901 return elif op in funcs_2args: if np.isscalar(value1) and np.isscalar(value2): - self.expression = f"{op}({value1}, {value2})" + self.expression = "o0" + self.operands = {"o0": ne_evaluate(f"{op}({value1}, {value2})")} # eager evaluation elif np.isscalar(value2): self.operands = {"o0": value1} self.expression = f"{op}(o0, {value2})" @@ -2288,7 +2311,8 @@ def __init__(self, new_op): # noqa: C901 self._dtype = dtype_ if np.isscalar(value1) and np.isscalar(value2): - self.expression = f"({value1} {op} {value2})" + self.expression = "o0" + self.operands = {"o0": ne_evaluate(f"({value1} {op} {value2})")} # eager evaluation elif np.isscalar(value2): self.operands = {"o0": value1} self.expression = f"(o0 {op} {value2})" @@ -2530,7 +2554,7 @@ def where(self, value1=None, value2=None): # This just acts as a 'decorator' for the existing expression if value1 is not None and value2 is not None: # Guess the outcome dtype for value1 and value2 - dtype = np.result_type(value1, value2) + dtype = blosc2.result_type(value1, value2) args = {"_where_x": value1, "_where_y": value2} elif value1 is not None: if hasattr(value1, "dtype"): @@ -2736,7 +2760,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 @@ -2881,7 +2905,7 @@ def __getitem__(self, item): # Squeeze single-element dimensions when indexing with integers # See e.g. examples/ndarray/animated_plot.py if isinstance(item, int) or (hasattr(item, "__iter__") and any(isinstance(i, int) for i in item)): - result = result.squeeze() + result = result.squeeze(axis=tuple(i for i in range(result.ndim) if result.shape[i] == 1)) return result def slice(self, item): @@ -3008,7 +3032,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()} diff --git a/src/blosc2/linalg.py b/src/blosc2/linalg.py index 5b56a5c2..63d7b887 100644 --- a/src/blosc2/linalg.py +++ b/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.as_simpleproxy(x1, x2) # Validate matrix multiplication compatibility if x1.shape[builtins.max(-1, -len(x2.shape))] != x2.shape[builtins.max(-2, -len(x2.shape))]: @@ -100,7 +99,7 @@ def matmul(x1: blosc2.Array, x2: blosc2.NDArray, **kwargs: Any) -> blosc2.NDArra n, k = x1.shape[-2:] m = x2.shape[-1] result_shape = np.broadcast_shapes(x1.shape[:-2], x2.shape[:-2]) + (n, m) - result = blosc2.zeros(result_shape, dtype=np.result_type(x1, x2), **kwargs) + result = blosc2.zeros(result_shape, dtype=blosc2.result_type(x1, x2), **kwargs) if 0 not in result.shape + x1.shape + x2.shape: # if any array is empty, return array of 0s p, q = result.chunks[-2:] @@ -183,11 +182,9 @@ 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) + # Makes a SimpleProxy if inputs are not blosc2 arrays + x1, x2 = blosc2.as_simpleproxy(x1, x2) if isinstance(axes, tuple): a_axes, b_axes = axes @@ -227,7 +224,7 @@ def tensordot( raise ValueError("x1 and x2 must have same shapes along reduction dimensions") result_shape = tuple(x1shape[a_keep]) + tuple(x2shape[b_keep]) - result = blosc2.zeros(result_shape, dtype=np.result_type(x1, x2), **kwargs) + result = blosc2.zeros(result_shape, dtype=blosc2.result_type(x1, x2), **kwargs) op_chunks = [ slice_to_chunktuple(slice(0, s, 1), c) for s, c in zip(x1shape[a_axes], a_chunks_red, strict=True) @@ -261,24 +258,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 +274,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 @@ -342,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, x2) N = builtins.min(x1.ndim, x2.ndim) if axis < -N or axis > -1: @@ -363,7 +347,7 @@ def vecdot(x1: blosc2.NDArray, x2: blosc2.NDArray, axis: int = -1, **kwargs) -> raise ValueError("x1 and x2 must have same shapes along reduction dimensions") result_shape = np.broadcast_shapes(x1shape[a_keep], x2shape[b_keep]) - result = blosc2.zeros(result_shape, dtype=np.result_type(x1, x2), **kwargs) + result = blosc2.zeros(result_shape, dtype=blosc2.result_type(x1, x2), **kwargs) res_chunks = [ slice_to_chunktuple(s, c) @@ -396,19 +380,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 @@ -486,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 @@ -506,9 +489,15 @@ def permute_dims( chunks = arr.chunks shape = arr.shape - - for info in arr.iterchunks_info(): - coords = info.coords + # handle SimpleProxy which doesn't have iterchunks_info + if hasattr(arr, "iterchunks_info"): + my_it = arr.iterchunks_info() + _get_el = lambda x: x.coords # noqa: E731 + else: + my_it = get_intersecting_chunks((), shape, chunks) + _get_el = lambda x: x.raw # noqa: E731 + for info in my_it: + coords = _get_el(info) start_stop = [ (coord * chunk, builtins.min(chunk * (coord + 1), dim)) for coord, chunk, dim in zip(coords, chunks, shape, strict=False) @@ -517,7 +506,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 @@ -555,7 +544,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.") @@ -579,6 +569,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] @@ -612,6 +604,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: @@ -648,6 +642,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.as_simpleproxy(x1, 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 diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index 23ef1168..c2422da6 100644 --- a/src/blosc2/ndarray.py +++ b/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.") @@ -4639,7 +4642,7 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray: return ndslice - def squeeze(self, axis=None) -> NDArray: + def squeeze(self, axis: int | Sequence[int]) -> NDArray: """Remove single-dimensional entries from the shape of the array. This method modifies the array in-place. If mask is None removes any dimensions with size 1. @@ -4663,18 +4666,15 @@ def squeeze(self, axis=None) -> NDArray: >>> a.shape (23, 11) """ - if axis is None: - super().squeeze() - else: - axis = [axis] if isinstance(axis, int) else axis - mask = [False for i in range(self.ndim)] - for a in axis: - if a < 0: - a += self.ndim # Adjust axis to be within the array's dimensions - if mask[a]: - raise ValueError("Axis values must be unique.") - mask[a] = True - super().squeeze(mask=mask) + axis = [axis] if isinstance(axis, int) else axis + mask = [False for i in range(self.ndim)] + for a in axis: + if a < 0: + a += self.ndim # Adjust axis to be within the array's dimensions + if mask[a]: + raise ValueError("Axis values must be unique.") + mask[a] = True + super().squeeze(mask=mask) return self def indices(self, order: str | list[str] | None = None, **kwargs: Any) -> NDArray: @@ -4708,17 +4708,23 @@ def __matmul__(self, other): return blosc2.linalg.matmul(self, other) -def squeeze(x: NDArray, axis: int | None = None) -> NDArray: +def squeeze(x: Array, axis: int | Sequence[int]) -> NDArray: """ Remove single-dimensional entries from the shape of the array. - This method modifies the array in-place. If mask is None removes any dimensions with size 1. - If axis is provided, it should be an int or tuple of ints and the corresponding - dimensions (of size 1) will be removed. + This method modifies the array in-place. + + Parameters + ---------- + x: Array + input array. + axis: int | Sequence[int] + Axis (or axes) to squeeze. Returns ------- - out: NDArray + out: Array + An output array having the same data type and elements as x. Examples -------- @@ -5653,7 +5659,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 +5671,14 @@ 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 +5689,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 diff --git a/src/blosc2/proxy.py b/src/blosc2/proxy.py index 13223087..f9fe5688 100644 --- a/src/blosc2/proxy.py +++ b/src/blosc2/proxy.py @@ -8,6 +8,12 @@ from abc import ABC, abstractmethod from collections.abc import Sequence +try: + from numpy.typing import DTypeLike +except (ImportError, AttributeError): + # fallback to internal module (use with caution) + from numpy._typing import DTypeLike + import numpy as np import blosc2 @@ -569,6 +575,20 @@ def __getitem__(self, item: slice | list[slice]) -> np.ndarray: return nparr[self.field] +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 + return np.dtype(str(dt).split(".")[1]) + except Exception as e: + raise TypeError("Could not parse dtype arg {dt}.") from e + + class SimpleProxy(blosc2.Operand): """ Simple proxy for any data container to be used with the compute engine. @@ -597,8 +617,8 @@ def __init__(self, src, chunks: tuple | None = None, blocks: tuple | None = None if not hasattr(src, "__getitem__"): raise TypeError("The source must have a __getitem__ method") self._src = src - self._dtype = src.dtype - self._shape = src.shape + self._dtype = _convert_dtype(src.dtype) + 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) @@ -629,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). @@ -642,7 +667,35 @@ def __getitem__(self, item: slice | list[slice]) -> np.ndarray: out: numpy.ndarray An array with the data slice. """ - return self._src[item] + out = self._src[item] + if not hasattr(out, "shape") or out.shape == (): + return out + else: + return np.asarray(out) # avoids copy for PyTorch at least + + +def as_simpleproxy(*arrs: Sequence[blosc2.Array]) -> tuple[SimpleProxy | blosc2.Operand]: + """ + Convert an Array object which fulfills Array protocol into SimpleProxy. If x is already a + blosc2.Operand simply returns object. + + Parameters + ---------- + arrs: Sequence[blosc2.Array] + Objects fulfilling Array protocol. + + Returns + ------- + out: tuple[blosc2.SimpleProxy | blosc2.Operand] + Objects with minimal interface for blosc2 LazyExpr computations. + """ + out = () + for x in arrs: + if isinstance(x, blosc2.Operand): + out += (x,) + else: + out += (SimpleProxy(x),) + return out[0] if len(out) == 1 else out def jit(func=None, *, out=None, disable=False, **kwargs): # noqa: C901 diff --git a/src/blosc2/shape_utils.py b/src/blosc2/shape_utils.py index 51c4122a..e278fa04 100644 --- a/src/blosc2/shape_utils.py +++ b/src/blosc2/shape_utils.py @@ -391,14 +391,7 @@ def visit_Call(self, node): # noqa : C901 # --- Parse keyword args --- kwargs = {} for kw in node.keywords: - if isinstance(kw.value, ast.Constant): - kwargs[kw.arg] = kw.value.value - elif isinstance(kw.value, ast.Tuple): - kwargs[kw.arg] = tuple( - e.value if isinstance(e, ast.Constant) else self._lookup_value(e) for e in kw.value.elts - ) - else: - kwargs[kw.arg] = self._lookup_value(kw.value) + kwargs[kw.arg] = self._lookup_value(kw.value) # ------- handle linear algebra --------------- if base_name in linalg_funcs: @@ -520,6 +513,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( @@ -536,15 +532,62 @@ def _eval_slice(self, node): else: raise ValueError(f"Unsupported slice expression: {ast.dump(node)}") - def _lookup_value(self, node): + def _lookup_value(self, node): # noqa : C901 """Look up a value in self.shapes if node is a variable name, else constant value.""" + # Name -> lookup in shapes mapping if isinstance(node, ast.Name): return self.shapes.get(node.id, None) - elif isinstance(node, ast.Constant): + + # Constant -> return its value + if isinstance(node, ast.Constant): return node.value - else: + + # Tuple of constants / expressions + if isinstance(node, ast.Tuple): + vals = [] + for e in node.elts: + v = self._lookup_value(e) + vals.append(v) + return tuple(vals) + + # Unary operations (e.g. -1) + if isinstance(node, ast.UnaryOp): + # handle negative constants like -1 + if isinstance(node.op, ast.USub): + val = self._lookup_value(node.operand) + if isinstance(val, (int, float)): + return -val + # handle + (USub) if needed + if isinstance(node.op, ast.UAdd): + return self._lookup_value(node.operand) + return None + + # Simple binary ops with constant operands (e.g. 1+2) + if isinstance(node, ast.BinOp): + left = self._lookup_value(node.left) + right = self._lookup_value(node.right) + if left is None or right is None: + return None + try: + if isinstance(node.op, ast.Add): + return left + right + if isinstance(node.op, ast.Sub): + return left - right + if isinstance(node.op, ast.Mult): + return left * right + if isinstance(node.op, ast.FloorDiv): + return left // right + if isinstance(node.op, ast.Div): + return left / right + if isinstance(node.op, ast.Mod): + return left % right + except Exception: + return None return None + # fallback + return None + # --- Public API --- def infer_shape(expr, shapes): diff --git a/tests/ndarray/test_elementwise_funcs.py b/tests/ndarray/test_elementwise_funcs.py index dfe99b9b..c6f95733 100644 --- a/tests/ndarray/test_elementwise_funcs.py +++ b/tests/ndarray/test_elementwise_funcs.py @@ -3,6 +3,7 @@ import numpy as np import pytest +import torch import blosc2 @@ -32,6 +33,7 @@ UNARY_FUNC_PAIRS.append((np.count_nonzero, blosc2.count_nonzero)) DTYPES = [blosc2.bool_, blosc2.int32, blosc2.int64, blosc2.float32, blosc2.float64, blosc2.complex128] +STR_DTYPES = ["bool", "int32", "int64", "float32", "float64", "complex128"] SHAPES_CHUNKS = [((10,), (3,)), ((20, 20), (4, 7))] SHAPES_CHUNKS_HEAVY = [((10, 13, 13), (3, 5, 2))] @@ -96,13 +98,132 @@ def _test_unary_func_impl(np_func, blosc_func, dtype, shape, chunkshape): raise e +def _test_binary_func_proxy(np_func, blosc_func, dtype, shape, chunkshape, xp): # noqa: C901 + dtype_ = getattr(xp, dtype) if hasattr(xp, dtype) else np.dtype(dtype) + dtype = np.dtype(dtype) + not_blosc1 = xp.ones(shape, dtype=dtype_) + if np_func.__name__ in ("right_shift", "left_shift"): + a_blosc2 = blosc2.asarray(2, copy=True) + else: + a_blosc2 = blosc2.linspace( + start=np.prod(shape) * 2, + stop=np.prod(shape), + num=np.prod(shape), + chunks=chunkshape, + shape=shape, + dtype=dtype, + ) + if not blosc2.isdtype(dtype, "integral"): + a_blosc2[tuple(i // 2 for i in shape)] = blosc2.nan + if dtype == blosc2.complex128: + a_blosc2 = ( + a_blosc2 + + blosc2.linspace( + 1j, + stop=np.prod(shape) * 1j, + num=np.prod(shape), + chunks=chunkshape, + shape=shape, + dtype=dtype, + ) + ).compute() + a_blosc2[tuple(i // 2 for i in shape)] = blosc2.nan + blosc2.nan * 1j + arr1 = np.asarray(not_blosc1) + arr2 = a_blosc2[()] + success = False + try: + expected = np_func(arr1, arr2) + success = True + except TypeError: + assert True + except RuntimeWarning as e: + assert True + if success: + try: + result = blosc_func(not_blosc1, a_blosc2)[()] + np.testing.assert_allclose(result, expected, rtol=1e-6, atol=1e-6) + except TypeError as e: + # some functions don't support certain dtypes and that's fine + assert True + except ValueError as e: # shouldn't be allowed for non-booleans + if np_func.__name__ in ("logical_and", "logical_or", "logical_xor"): + assert True + if ( + np_func.__name__ in ("less", "less_equal", "greater", "greater_equal", "minimum", "maximum") + and dtype == blosc2.complex128 + ): # not supported for complex dtypes + assert True + else: + raise e + except NotImplementedError as e: + if np_func.__name__ in ("left_shift", "right_shift", "floor_divide", "power", "remainder"): + assert True + else: + raise e + except AssertionError as e: + if np_func.__name__ == "power" and blosc2.isdtype( + dtype, "integral" + ): # overflow causes disagreement, no problem + assert True + elif np_func.__name__ in ("maximum", "minimum") and blosc2.isdtype(dtype, "real floating"): + warnings.showwarning( + "minimum and maximum for numexpr do not match NaN behaviour for numpy", + UserWarning, + __file__, + 0, + file=sys.stderr, + ) + pytest.skip("minimum and maximum for numexpr do not match NaN behaviour for numpy") + else: + raise e + + +def _test_unary_func_proxy(np_func, blosc_func, dtype, shape, xp): + dtype_ = getattr(xp, dtype) if hasattr(xp, dtype) else np.dtype(dtype) + dtype = np.dtype(dtype) + a_blosc = xp.ones(shape, dtype=dtype_) + if not blosc2.isdtype(dtype, "integral"): + a_blosc[tuple(i // 2 for i in shape)] = xp.nan + if dtype == blosc2.complex128: + 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, dtype=dtype_) + + arr = np.asarray(a_blosc) + success = False + try: + expected = np_func(arr) if np_func.__name__ != "reciprocal" else 1.0 / arr + success = True + except TypeError: + assert True + except RuntimeWarning as e: + assert True + if success: + try: + result = blosc_func(a_blosc)[...] + np.testing.assert_allclose(result, expected, rtol=1e-6, atol=1e-6) + except TypeError as e: + # some functions don't support certain dtypes and that's fine + assert True + except ValueError as e: + if np_func.__name__ == "logical_not" and dtype in ( + blosc2.float32, + blosc2.float64, + blosc2.complex128, + ): + assert True + else: + raise e + + def _test_binary_func_impl(np_func, blosc_func, dtype, shape, chunkshape): # noqa: C901 """Helper function containing the actual test logic for binary functions.""" a_blosc1 = blosc2.linspace( 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, @@ -179,6 +300,14 @@ def test_unary_funcs(np_func, blosc_func, dtype, shape, chunkshape): _test_unary_func_impl(np_func, blosc_func, dtype, shape, chunkshape) +@pytest.mark.parametrize(("np_func", "blosc_func"), UNARY_FUNC_PAIRS) +@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, xp): + _test_unary_func_proxy(np_func, blosc_func, dtype, shape, xp) + + @pytest.mark.heavy @pytest.mark.parametrize(("np_func", "blosc_func"), UNARY_FUNC_PAIRS) @pytest.mark.parametrize("dtype", DTYPES) @@ -194,6 +323,14 @@ def test_binary_funcs(np_func, blosc_func, dtype, shape, chunkshape): _test_binary_func_impl(np_func, blosc_func, dtype, shape, chunkshape) +@pytest.mark.parametrize(("np_func", "blosc_func"), BINARY_FUNC_PAIRS) +@pytest.mark.parametrize("dtype", STR_DTYPES) +@pytest.mark.parametrize(("shape", "chunkshape"), SHAPES_CHUNKS) +@pytest.mark.parametrize("xp", [torch]) +def test_binfuncs_proxy(np_func, blosc_func, dtype, shape, chunkshape, xp): + _test_binary_func_proxy(np_func, blosc_func, dtype, shape, chunkshape, xp) + + @pytest.mark.heavy @pytest.mark.parametrize(("np_func", "blosc_func"), BINARY_FUNC_PAIRS) @pytest.mark.parametrize("dtype", DTYPES) diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 5d6156d3..3ca6b473 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -10,6 +10,7 @@ import numpy as np import pytest +import torch import blosc2 from blosc2.lazyexpr import ne_evaluate @@ -358,12 +359,7 @@ def test_functions(function, dtype_fixture, shape_fixture): ) @pytest.mark.parametrize( ("value1", "value2"), - [ - ("NDArray", "scalar"), - ("NDArray", "NDArray"), - ("scalar", "NDArray"), - # ("scalar", "scalar") # Not supported by LazyExpr - ], + [("NDArray", "scalar"), ("NDArray", "NDArray"), ("scalar", "NDArray"), ("scalar", "scalar")], ) def test_arctan2_pow(urlpath, shape_fixture, dtype_fixture, function, value1, value2): nelems = np.prod(shape_fixture) @@ -405,7 +401,7 @@ def test_arctan2_pow(urlpath, shape_fixture, dtype_fixture, function, value1, va else: expr_string = f"{function}(na1, value2)" res_numexpr = ne_evaluate(expr_string) - else: # ("scalar", "NDArray") + elif value2 == "NDArray": # ("scalar", "NDArray") value1 = 12 na2 = np.linspace(0, 10, nelems, dtype=dtype_fixture).reshape(shape_fixture) a2 = blosc2.asarray(na2, urlpath=urlpath2, mode="w") @@ -421,9 +417,21 @@ def test_arctan2_pow(urlpath, shape_fixture, dtype_fixture, function, value1, va else: expr_string = f"{function}(value1, na2)" res_numexpr = ne_evaluate(expr_string) + else: # ("scalar", "scalar") + value1 = 12 + value2 = 3 + # Construct the lazy expression based on the function name + expr = blosc2.LazyExpr(new_op=(value1, function, value2)) + res_lazyexpr = expr.compute() + # Evaluate using NumExpr + if function == "**": + res_numexpr = ne_evaluate("value1**value2") + else: + expr_string = f"{function}(value1, value2)" + res_numexpr = ne_evaluate(expr_string) # Compare the results tol = 1e-15 if dtype_fixture == "float64" else 1e-6 - np.testing.assert_allclose(res_lazyexpr[:], res_numexpr, atol=tol, rtol=tol) + np.testing.assert_allclose(res_lazyexpr[()], res_numexpr, atol=tol, rtol=tol) for path in [urlpath1, urlpath2, urlpath_save]: blosc2.remove_urlpath(path) @@ -1195,7 +1203,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) @@ -1510,6 +1518,11 @@ def test_scalar_dtypes(values): dtype2 = (avalue1 * avalue2).dtype assert dtype1 == dtype2, f"Expected {dtype1} but got {dtype2}" + # test scalars + value = value1 if np.isscalar(value1) else value2 + assert blosc2.sin(value)[()] == np.sin(value) + assert (value + blosc2.sin(value))[()] == value + np.sin(value) + def test_to_cframe(): N = 1_000 @@ -1697,13 +1710,15 @@ def test_lazylinalg(): np.testing.assert_array_almost_equal(out[()], npres) # --- squeeze --- - out = blosc2.lazyexpr("squeeze(D)") - npres = np.squeeze(npD) + out = blosc2.lazyexpr("squeeze(D, axis=-1)") + npres = np.squeeze(npD, -1) assert out.shape == npres.shape np.testing.assert_array_almost_equal(out[()], npres) - - out = blosc2.lazyexpr("D.squeeze()") - npres = np.squeeze(npD) + # refresh D since squeeze is in-place + s = shapes["D"] + D = blosc2.linspace(0, np.prod(s), shape=s) + out = blosc2.lazyexpr("D.squeeze(axis=-1)") + npres = np.squeeze(npD, -1) assert out.shape == npres.shape np.testing.assert_array_almost_equal(out[()], npres) @@ -1714,10 +1729,14 @@ def test_lazylinalg(): np.testing.assert_array_almost_equal(out[()], npres) # --- tensordot --- - out = blosc2.lazyexpr("tensordot(A, B, axes=1)") + out = blosc2.lazyexpr("tensordot(A, B, axes=1)") # test with int axes npres = np.tensordot(npA, npB, axes=1) assert out.shape == npres.shape np.testing.assert_array_almost_equal(out[()], npres) + out = blosc2.lazyexpr("tensordot(A, B, axes=((1,) , (0,)))") # test with tuple axes + npres = np.tensordot(npA, npB, axes=((1,), (0,))) + assert out.shape == npres.shape + np.testing.assert_array_almost_equal(out[()], npres) # --- vecdot --- out = blosc2.lazyexpr("vecdot(x, y)") @@ -1768,3 +1787,55 @@ def test_lazyexpr_2args(): newexpr = blosc2.hypot(lexpr, 3) assert newexpr.expression == "hypot((sin(o0)), 3)" assert newexpr.operands["o0"] is a + + +@pytest.mark.parametrize( + "xp", + [torch, np], +) +@pytest.mark.parametrize( + "dtype", + ["bool", "int32", "int64", "float32", "float64", "complex128"], +) +def test_simpleproxy(xp, dtype): + try: + dtype_ = getattr(xp, dtype) if hasattr(xp, dtype) else np.dtype(dtype) + except FutureWarning: + dtype_ = 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) diff --git a/tests/ndarray/test_linalg.py b/tests/ndarray/test_linalg.py index fd6ecb0f..988f00f8 100644 --- a/tests/ndarray/test_linalg.py +++ b/tests/ndarray/test_linalg.py @@ -1,9 +1,12 @@ +import inspect from itertools import permutations import numpy as np import pytest +import torch import blosc2 +from blosc2.lazyexpr import linalg_funcs from blosc2.ndarray import npvecdot @@ -817,3 +820,64 @@ def test_diagonal(shape, chunkshape, offset): # Assert equality np.testing.assert_array_equal(result_np, expected) + + +@pytest.mark.parametrize( + "xp", + [torch, np], +) +@pytest.mark.parametrize( + "dtype", + ["int32", "int64", "float32", "float64", "complex128"], +) +def test_linalgproxy(xp, dtype): + dtype_ = getattr(xp, dtype) if hasattr(xp, dtype) else np.dtype(dtype) + for name in linalg_funcs: + if name == "transpose": + continue # deprecated + func = getattr(blosc2, name) + N = 10 + shape_a = (N,) + chunks = (N // 3,) + if name != "outer": + shape_a *= 3 + chunks *= 3 + blosc_matrix = blosc2.full(shape=shape_a, fill_value=3, dtype=np.dtype(dtype), chunks=chunks) + 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=chunks + ) + + # Check this works + argspec = inspect.getfullargspec(func) + num_args = len(argspec.args) + # handle numpy 1.26 + if name == "permute_dims": + npfunc = blosc2.linalg.nptranspose + elif name == "concat" and not hasattr(np, "concat"): + npfunc = np.concatenate + elif name == "matrix_transpose": + npfunc = blosc2.linalg.nptranspose + elif name == "vecdot": + npfunc = blosc2.linalg.npvecdot + else: + npfunc = getattr(np, name) + if num_args > 2 or name in ("outer", "matmul"): + try: + lexpr = func(blosc_matrix, foreign_matrix) + except NotImplementedError: + continue + foreign_matrix = np.asarray(foreign_matrix) + res = npfunc(blosc_matrix[()], foreign_matrix) + else: + try: + lexpr = func(foreign_matrix) + except NotImplementedError: + continue + except TypeError: + continue + foreign_matrix = np.asarray(foreign_matrix) + res = npfunc(foreign_matrix, 0) if name == "expand_dims" else npfunc(foreign_matrix) + np.testing.assert_array_equal(res, lexpr[()]) diff --git a/tests/ndarray/test_squeeze.py b/tests/ndarray/test_squeeze.py index e10b871e..397e5351 100644 --- a/tests/ndarray/test_squeeze.py +++ b/tests/ndarray/test_squeeze.py @@ -19,7 +19,6 @@ ((23, 1, 1, 34), (20, 1, 1, 20), None, 1234, 2), ((80, 1, 51, 60, 1), None, (6, 1, 6, 26, 1), 3.333, 4), ((1, 1, 1), None, None, True, (1, 2)), - ((1, 1, 1), None, None, True, None), ], ) def test_squeeze(shape, chunks, blocks, fill_value, axis):