From d153173afd0b73ff984f4348437da2c86228dcb5 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Tue, 3 Nov 2026 13:19:10 +0100 Subject: [PATCH 1/3] Keep std/var and where() with scalars on the miniexpr path Two common cases fell off miniexpr, and without numexpr they now run chunk by chunk through evaluate_chunk(): - Full std()/var() of float32/float64 (no axis, no where, same dtype) reduced (x - mean) ** 2 with the mean as a broadcast (1, ..., 1) operand. They are now two single-pass reductions: the mean, then the mean of (x - m) * (x - m) with m inlined as a scalar. The sign of m is folded into the operator, as miniexpr reductions mis-evaluate a negative literal on the right of a binary operator. Other cases keep the general path; std() and var() share one implementation. - where(cond, x, y) with a scalar x or y: specialize_miniexpr_inputs() inlined the scalar and dropped its operand before the where() was fused in, so the fused expression referenced a missing operand, miniexpr failed to compile it and the evaluation fell back. The where() is now fused before specializing. - The bare-array LazyExpr that the NDArray reductions build takes the array's dtype instead of inferring it. Measured on a Ryzen 9 7950X (300x512x512 float32 in memory, 28 threads), against this branch at 7db7b09: a.std() 470 -> 53 ms, (a.mean(), a.std()) 498 -> 78 ms, where(a > 100, a, 0).compute() 690 -> 469 ms; 100x100 a.std() 0.73 -> 0.52 ms. std/var may differ from before in the last ulps. Taken from the lazy-expression commit of #744 (b28fb0e), which was deferred until this branch lands. Co-Authored-By: Claude Opus 5.5 --- src/blosc2/lazyexpr.py | 185 +++++++++++++++---------------- tests/ndarray/test_lazyexpr.py | 40 +++++++ tests/ndarray/test_reductions.py | 46 ++++++++ 3 files changed, 176 insertions(+), 95 deletions(-) diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 2229cf7f2..52a01e0ff 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -1887,13 +1887,23 @@ def fast_eval( # noqa: C901 expr_string_miniexpr = expr_string operands_miniexpr = operands + # Fuse where(cond, x, y) before specializing, so that scalar x/y get inlined as well; + # fusing afterwards would reference a dropped operand and make miniexpr fail to compile. + where_fused = ( + use_miniexpr + and isinstance(expr_string, str) + and not is_dsl + and where is not None + and len(where) == 2 + ) + if where_fused: + expr_string_miniexpr = f"where({expr_string}, _where_x, _where_y)" if use_miniexpr and isinstance(expr_string, str): - try: - expr_string_miniexpr, operands_miniexpr = specialize_miniexpr_inputs(expr_string, operands) - except Exception: - # If specialization fails, keep original expression/operands and let normal checks decide. - expr_string_miniexpr = expr_string - operands_miniexpr = operands + # If specialization fails, keep original expression/operands and let normal checks decide. + with contextlib.suppress(Exception): + expr_string_miniexpr, operands_miniexpr = specialize_miniexpr_inputs( + expr_string_miniexpr, operands + ) # Check whether we can use miniexpr if use_miniexpr: @@ -2025,7 +2035,7 @@ def _miniexpr_eligible_operand(op): k: (np.ascontiguousarray(v) if isinstance(v, np.ndarray) else v) for k, v in _pref_ops.items() } - if where is not None and len(where) == 2: + if where is not None and len(where) == 2 and not where_fused: _pref_expr = f"where({_pref_expr}, _where_x, _where_y)" # _cb_anchor keeps the contiguous bitmap alive for the whole # update_data loop (the prefilter dereferences it per block). Pass @@ -3774,6 +3784,9 @@ def __init__(self, new_op): # noqa: C901 op = None self.operands = {"o0": value1} self.expression = "o0" if op is None else f"{op}(o0)" + if op is None and isinstance(value1, blosc2.NDArray): + # A bare array (as the NDArray reductions build): no need to infer the dtype + self._dtype = value1.dtype return value2 = ( blosc2.SimpleProxy(value2) @@ -4301,51 +4314,69 @@ def mean( out = blosc2.asarray(out, **kwargs) return out - @expression_method - def std( - self, - axis=None, - dtype=None, - keepdims=False, - ddof=0, - where=None, - fp_accuracy: blosc2.FPAccuracy | None = None, - **kwargs, - ): + def _can_full_var(self, dtype): + """Whether _full_var computes the variance in `dtype`: for float32/float64 operands, where + miniexpr evaluates it without per-element conversions, into their own dtype (the inlined + mean would not make the deviations of float32 operands float64).""" + own = np.dtype(self.dtype) + return own.kind == "f" and own.itemsize in (4, 8) and (dtype is None or np.dtype(dtype) == own) + + def _full_var(self, dtype, keepdims, fp_accuracy, execution_kwargs): + """Variance (ddof=0) over all the elements, as two full reductions over the operands. + + The mean goes into the second pass as a Python scalar, which gets inlined into the + expression; a (1, ..., 1) array would be an operand to broadcast, and keep that reduction + off its fast paths. d * d rather than d ** 2 spares a pow() per element. Only for + float32/float64 into the same dtype (see _can_full_var). + """ + mean_value = np.asarray(self.mean(dtype=dtype, fp_accuracy=fp_accuracy, **execution_kwargs)).item() + # Fold the sign into the operator: miniexpr reductions mis-evaluate a negative literal on + # the right of a binary operator, as in "(o0 - -0.5)" + dev = self + (-mean_value) if math.copysign(1.0, mean_value) < 0 else self - mean_value + out = (dev * dev).mean(dtype=dtype, fp_accuracy=fp_accuracy, **execution_kwargs) + # Evaluators other than miniexpr may compute against the Python scalar in float64 + out = np.dtype(self.dtype if dtype is None else dtype).type(out) + return np.reshape(out, (1,) * self.ndim) if keepdims else out + + def _var(self, axis, dtype, keepdims, ddof, where, fp_accuracy, kwargs, finish=None): + """The variance, as var() returns it; std() passes np.sqrt as `finish`.""" execution_kwargs = blosc2.jit_config.pop_execution_options(kwargs) where = self._normalize_where(where) item = kwargs.pop("item", ()) - if item == (): # fast path - mean_value = self.mean( - axis=axis, - dtype=dtype, - keepdims=True, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - expr = (self - mean_value) ** 2 + if item == () and axis is None and where is None and self._can_full_var(dtype): + out = self._full_var(dtype, keepdims, fp_accuracy, execution_kwargs) else: - mean_value = self.mean( + if item == (): # fast path + mean_value = self.mean( + axis=axis, + dtype=dtype, + keepdims=True, + where=where, + fp_accuracy=fp_accuracy, + **execution_kwargs, + ) + expr = (self - mean_value) ** 2 + else: + mean_value = self.mean( + axis=axis, + dtype=dtype, + keepdims=True, + where=where, + item=item, + fp_accuracy=fp_accuracy, + **execution_kwargs, + ) + # TODO: Not optimal because we load the whole slice in memory. Would have to write + # a bespoke var function that executed within slice_eval to avoid this probably. + expr = (self.compute(item, fp_accuracy=fp_accuracy, **execution_kwargs) - mean_value) ** 2 + out = expr.mean( axis=axis, dtype=dtype, - keepdims=True, + keepdims=keepdims, where=where, - item=item, fp_accuracy=fp_accuracy, **execution_kwargs, ) - # TODO: Not optimal because we load the whole slice in memory. Would have to write - # a bespoke std function that executed within slice_eval to avoid this probably. - expr = (self.compute(item, fp_accuracy=fp_accuracy, **execution_kwargs) - mean_value) ** 2 - out = expr.mean( - axis=axis, - dtype=dtype, - keepdims=keepdims, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) if ddof != 0: num_elements = ( self.get_num_elements(axis, item, fp_accuracy=fp_accuracy, **execution_kwargs) @@ -4354,9 +4385,9 @@ def std( axis=axis, dtype=np.int64, fp_accuracy=fp_accuracy, **execution_kwargs ) ) - out = np.sqrt(out * num_elements / (num_elements - ddof)) - else: - out = np.sqrt(out) + out = out * num_elements / (num_elements - ddof) + if finish is not None: + out = finish(out) out2 = kwargs.pop("out", None) if out2 is not None: out2[:] = out @@ -4365,6 +4396,19 @@ def std( out = blosc2.asarray(out, **kwargs) return out + @expression_method + def std( + self, + axis=None, + dtype=None, + keepdims=False, + ddof=0, + where=None, + fp_accuracy: blosc2.FPAccuracy | None = None, + **kwargs, + ): + return self._var(axis, dtype, keepdims, ddof, where, fp_accuracy, kwargs, np.sqrt) + @expression_method def var( self, @@ -4376,56 +4420,7 @@ def var( fp_accuracy: blosc2.FPAccuracy | None = None, **kwargs, ): - execution_kwargs = blosc2.jit_config.pop_execution_options(kwargs) - where = self._normalize_where(where) - item = kwargs.pop("item", ()) - if item == (): # fast path - mean_value = self.mean( - axis=axis, - dtype=dtype, - keepdims=True, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - expr = (self - mean_value) ** 2 - else: - mean_value = self.mean( - axis=axis, - dtype=dtype, - keepdims=True, - where=where, - item=item, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - # TODO: Not optimal because we load the whole slice in memory. Would have to write - # a bespoke var function that executed within slice_eval to avoid this probably. - expr = (self.compute(item, fp_accuracy=fp_accuracy, **execution_kwargs) - mean_value) ** 2 - out = expr.mean( - axis=axis, - dtype=dtype, - keepdims=keepdims, - where=where, - fp_accuracy=fp_accuracy, - **execution_kwargs, - ) - if ddof != 0: - num_elements = ( - self.get_num_elements(axis, item, fp_accuracy=fp_accuracy, **execution_kwargs) - if where is None - else where.where(blosc2.ones(self.shape, dtype=np.int64), 0).sum( - axis=axis, dtype=np.int64, fp_accuracy=fp_accuracy, **execution_kwargs - ) - ) - out = out * num_elements / (num_elements - ddof) - out2 = kwargs.pop("out", None) - if out2 is not None: - out2[:] = out - return out2 - if kwargs != {} and not np.isscalar(out): - out = blosc2.asarray(out, **kwargs) - return out + return self._var(axis, dtype, keepdims, ddof, where, fp_accuracy, kwargs) @expression_method def min( diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 4be24af5d..7a1223e5e 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -2436,3 +2436,43 @@ def test_cumulative_block_layout(monkeypatch, blocks): expected = np.cumsum(data, axis=0) np.testing.assert_array_equal(result[:], expected) np.testing.assert_array_equal(out[:], expected) + + +@pytest.mark.parametrize( + ("dtype", "x", "y"), + [ + (np.float32, "a", 0), + (np.float32, 1.5, "a"), + (np.int32, "a", 0.5), + (np.int16, "a", -1), + (np.float64, "a", "b"), + ], +) +@pytest.mark.parametrize("getitem", [False, True]) +def test_where_scalar_uses_miniexpr(monkeypatch, dtype, x, y, getitem): + import importlib + + lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") + monkeypatch.setattr(lazyexpr_mod, "try_miniexpr", True) + original_set_pref_expr = blosc2.NDArray._set_pref_expr + compiled = [] + + def wrapped_set_pref_expr(self, expression, inputs, *args, **kwargs): + res = original_set_pref_expr(self, expression, inputs, *args, **kwargs) + compiled.append(expression) + return res + + monkeypatch.setattr(blosc2.NDArray, "_set_pref_expr", wrapped_set_pref_expr) + na = np.arange(40 * 30, dtype=dtype).reshape(40, 30) + nb = (na // 3).astype(dtype) + a = blosc2.asarray(na, chunks=(16, 30), blocks=(4, 30)) + b = blosc2.asarray(nb, chunks=(16, 30), blocks=(4, 30)) + ops, nops = {"a": a, "b": b}, {"a": na, "b": nb} + expr = blosc2.where(a > 500, ops.get(x, x), ops.get(y, y)) + expected = np.where(na > 500, nops.get(x, x), nops.get(y, y)) + res = expr[:] if getitem else expr.compute()[:] + assert res.dtype == expected.dtype + np.testing.assert_array_equal(res, expected) + # The scalar is inlined, so miniexpr (rather than the numexpr fallback) did the work + assert len(compiled) == 1 + assert "_where_x" in compiled[0] or "_where_y" in compiled[0] diff --git a/tests/ndarray/test_reductions.py b/tests/ndarray/test_reductions.py index df4951068..5c11ef8de 100644 --- a/tests/ndarray/test_reductions.py +++ b/tests/ndarray/test_reductions.py @@ -983,3 +983,49 @@ def test_reductions_where_ndarray_and_lazyexpr(): np.testing.assert_allclose(a.std(axis=0, where=mask), np.std(na, axis=0, where=nmask)) np.testing.assert_allclose(a.min(axis=0, where=mask), np.min(na, axis=0, where=nmask, initial=np.inf)) np.testing.assert_allclose(a.max(axis=0, where=mask), np.max(na, axis=0, where=nmask, initial=-np.inf)) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("op", ["std", "var"]) +@pytest.mark.parametrize("ddof", [0, 1]) +@pytest.mark.parametrize("keepdims", [False, True]) +@pytest.mark.parametrize("offset", [0, 1e4, -1e4]) +@pytest.mark.parametrize("miniexpr", [True, False]) +def test_std_var_full_reduction(monkeypatch, dtype, op, ddof, keepdims, offset, miniexpr): + # Full std/var reduce (x - mean) * (x - mean) with the mean inlined as a scalar + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", miniexpr) + rng = np.random.default_rng(42) + na = (offset + rng.normal(scale=3, size=(37, 40, 50))).astype(dtype) + a = blosc2.asarray(na, chunks=(10, 20, 50), blocks=(5, 10, 25)) + res = getattr(a, op)(ddof=ddof, keepdims=keepdims) + expected = getattr(na.astype(np.float64), op)(ddof=ddof, keepdims=keepdims) + assert np.asarray(res).dtype == getattr(na, op)().dtype + assert np.shape(res) == np.shape(expected) + np.testing.assert_allclose(res, expected, rtol=1e-5 if dtype == np.float32 else 1e-12) + # Same through a lazy expression and the module-level functions + res = getattr(blosc2, op)(a + 1, ddof=ddof, keepdims=keepdims) + assert np.asarray(res).dtype == getattr(na, op)().dtype + np.testing.assert_allclose(res, expected, rtol=1e-5 if dtype == np.float32 else 1e-12) + + +@pytest.mark.parametrize("op", ["std", "var"]) +def test_std_var_full_reduction_wider_dtype(op): + # A float64 result of float32 data takes float64 deviations from the mean, as NumPy does + rng = np.random.default_rng(1) + na = (rng.random(400_000) + 1e6).astype(np.float32) + a = blosc2.asarray(na) + res = getattr(a, op)(dtype=np.float64) + assert np.asarray(res).dtype == np.float64 + np.testing.assert_allclose(res, getattr(np, op)(na, dtype=np.float64), rtol=1e-6) + + +@pytest.mark.parametrize("value", [np.nan, np.inf]) +def test_std_var_full_reduction_nonfinite(value): + na = np.arange(1000, dtype=np.float64).reshape(10, 100) + na[3, 7] = value + a = blosc2.asarray(na, chunks=(5, 50), blocks=(5, 10)) + with np.errstate(invalid="ignore"): + np.testing.assert_array_equal(a.std(), np.std(na)) + np.testing.assert_array_equal(a.var(), np.var(na)) From cb44e45787d9bd60cfc986e87497900e1f93cc29 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Tue, 3 Nov 2026 14:23:36 +0100 Subject: [PATCH 2/3] Fuse where() into miniexpr only where it computes as NumPy does Fusing where(cond, x, y) with scalar x/y (previous commit) exposed miniexpr to cases it gets wrong, as Francesc found in review: - where(a > 0.1, 1, 0) with a float32 `a` of 0.1 selected 1. miniexpr types the literals of the fused expression by its output dtype and computes in it: with an int64 output 0.1 became 0, with a float32 output it compared against the double 0.1, where NumPy rounds it to float32 first. Arithmetic in the condition also ran in the output dtype (where(a + b == c, 1, 0) with float32 operands). - where(a > 5, np.uint64(2**63 + 1), a) gave 2**63: miniexpr parses literals into C doubles. Inside where() it also passes 64-bit integer operands and branch values through doubles. where() is now fused only when computing the condition in the output dtype changes nothing: all its operands have the output dtype, or it only compares operands that cast safely to it against literals their dtypes hold exactly; never with 64-bit integers. Otherwise it falls back to the chunked evaluation. This also fixes array branches, which were fused before and had the same problems (where(a > 0.1, a, b)). Masks, label maps, thresholds and same-dtype arithmetic stay on miniexpr; where(a > 100, a, 0).compute() keeps its speed. Expressions with integer literals beyond 2**53, written or inlined scalars, are no longer sent to miniexpr (also a + np.uint64(2**63 + 1)). Tests: the review examples, a NumPy differential grid over condition dtypes, literals, comparisons, output dtypes and scalar/array branches, arithmetic conditions, large integer constants, and checks that the common mixed-dtype where() patterns still compile with miniexpr. Co-Authored-By: Claude Opus 5.5 --- src/blosc2/lazyexpr.py | 153 +++++++++++++++++++++++++ tests/ndarray/test_lazyexpr.py | 201 ++++++++++++++++++++++++++++++--- 2 files changed, 341 insertions(+), 13 deletions(-) diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 52a01e0ff..8e40e9e9b 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -21,6 +21,7 @@ import re import sys import threading +import warnings from abc import ABC, abstractmethod, abstractproperty from collections import OrderedDict from collections.abc import MutableMapping @@ -1766,6 +1767,153 @@ def _miniexpr_integer_atan2(expression, operands): ) +def _literal_value(node): + """The number a literal AST node stands for (signs included), or None.""" + sign = 1 + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub | ast.UAdd): + sign = -1 if isinstance(node.op, ast.USub) else 1 + node = node.operand + if isinstance(node, ast.Constant) and isinstance(node.value, bool | int | float): + return sign * node.value + return None + + +def _scalar_operand_value(value): + """The number a scalar operand stands for, or None for arrays (and anything else).""" + if hasattr(value, "shape"): + if value.shape != (): + return None + value = value[()] + if isinstance(value, np.generic): + value = value.item() + return value if isinstance(value, bool | int | float) else None + + +def _exact_in_dtype(value, dtype) -> bool: + """Whether `dtype`, and the C double miniexpr keeps literals in, hold the number exactly.""" + if dtype.kind not in "biuf" or float(value) != value: + return False + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return dtype.type(value).item() == value + except (OverflowError, ValueError): + return False + + +@lru_cache(maxsize=1024) +def _parse_expression(expression: str) -> ast.Expression | None: + """The AST of an expression (None if it is not Python syntax); the callers do not modify it.""" + try: + return ast.parse(expression, mode="eval") + except SyntaxError: + return None + + +@lru_cache(maxsize=1024) +def _miniexpr_inexact_int_literal(expression: str) -> bool: + """Whether the expression holds an integer literal beyond 2**53. + + miniexpr parses numeric literals into C doubles, which round such integers. + """ + tree = _parse_expression(expression) + return tree is not None and any( + isinstance(node, ast.Constant) + and isinstance(node.value, int) + and not isinstance(node.value, bool) + and float(node.value) != node.value + for node in ast.walk(tree) + ) + + +def _miniexpr_where_matches_numpy(cond: str, operands: dict, dtype) -> bool: # noqa: C901 + """Whether miniexpr computes where(cond, x, y) as NumPy does, once the where() is fused in. + + miniexpr types the literals of the fused expression by its output dtype, and computes in it, + while NumPy evaluates the condition in the dtypes of its operands, with Python scalars as weak + scalars (NEP 50). With a float32 `a`, NumPy compares ``a > 0.1`` against float32(0.1), but + miniexpr compares ``where(a > 0.1, 1, 0)`` against the double 0.1, and adds ``a + b`` in + ``where(a + b == c, 1, 0)`` in a wider type. The two agree when computing in the output dtype + changes nothing: when every operand of the condition has the output dtype, or when the + condition only compares operands that cast safely to it, against literals that their dtypes + hold exactly. Inside where(), miniexpr also passes 64-bit integers through doubles, so + where() on them stays with NumPy. + """ + if dtype is None: + return False + out_dtype = np.dtype(dtype) + dtypes = [out_dtype] + [ + np.dtype(op.dtype) + for op in operands.values() + if hasattr(op, "dtype") and _scalar_operand_value(op) is None + ] + if any(dt.kind in "iu" and dt.itemsize > 4 for dt in dtypes): + return False + tree = _parse_expression(cond) + if tree is None: + return False + + def term(node): + """("literal", number) or ("array", dtype) for a leaf of the condition, else None.""" + value = _literal_value(node) + if value is not None: + return "literal", value + if isinstance(node, ast.Name) and node.id in operands: + op = operands[node.id] + value = _scalar_operand_value(op) + if value is not None: + return "literal", value + if hasattr(op, "dtype"): + return "array", np.dtype(op.dtype) + return None + + comparisons = [] + + def only_compares(node): + """Whether `node` combines comparisons of leaves with & | ~ only (collected into comparisons).""" + if isinstance(node, ast.Compare): + terms = [term(n) for n in (node.left, *node.comparators)] + if len(node.ops) > 1 or None in terms: + return False + comparisons.append(terms) + return True + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitAnd | ast.BitOr | ast.BitXor): + return only_compares(node.left) and only_compares(node.right) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Invert | ast.Not): + return only_compares(node.operand) + if isinstance(node, ast.BoolOp): + return all(only_compares(value) for value in node.values) + leaf = term(node) + if leaf is not None and leaf[0] == "array" and leaf[1].kind == "b": + comparisons.append([leaf]) + return True + return False + + if only_compares(tree.body): + for terms in comparisons: + dtypes = [dt for kind, dt in terms if kind == "array"] + if not all(np.can_cast(dt, out_dtype, "safe") for dt in dtypes): + return False + literals = [value for kind, value in terms if kind == "literal"] + if not all(_exact_in_dtype(value, dt) for value in literals for dt in (*dtypes, out_dtype)): + return False + return True + + # With arithmetic, the condition is computed as in NumPy only in its own dtype + def same_dtype(node): + leaf = term(node) + if leaf is not None: + kind, value = leaf + return value == out_dtype if kind == "array" else _exact_in_dtype(value, out_dtype) + if isinstance(node, ast.Name): # a function name + return True + if isinstance(node, ast.Constant): # strings and other non-numbers + return False + return all(same_dtype(child) for child in ast.iter_child_nodes(node)) + + return same_dtype(tree.body) + + @blosc2.jit_config.jit_execution def fast_eval( # noqa: C901 expression: str | Callable[[tuple, np.ndarray, tuple[int]], None], @@ -1896,6 +2044,8 @@ def fast_eval( # noqa: C901 and where is not None and len(where) == 2 ) + if where_fused and not _miniexpr_where_matches_numpy(expr_string, operands, dtype): + use_miniexpr = where_fused = False if where_fused: expr_string_miniexpr = f"where({expr_string}, _where_x, _where_y)" if use_miniexpr and isinstance(expr_string, str): @@ -1904,6 +2054,9 @@ def fast_eval( # noqa: C901 expr_string_miniexpr, operands_miniexpr = specialize_miniexpr_inputs( expr_string_miniexpr, operands ) + if not is_dsl and _miniexpr_inexact_int_literal(expr_string_miniexpr): + # Typed literals and inlined scalars alike + use_miniexpr = False # Check whether we can use miniexpr if use_miniexpr: diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 7a1223e5e..c2e38bb3c 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -5,7 +5,9 @@ # SPDX-License-Identifier: BSD-3-Clause ####################################################################### +import itertools import math +import operator import pathlib import warnings @@ -2438,6 +2440,23 @@ def test_cumulative_block_layout(monkeypatch, blocks): np.testing.assert_array_equal(out[:], expected) +def _record_miniexpr(monkeypatch): + """Force miniexpr on, and return the list the expressions it compiles get appended to.""" + import importlib + + monkeypatch.setattr(importlib.import_module("blosc2.lazyexpr"), "try_miniexpr", True) + original_set_pref_expr = blosc2.NDArray._set_pref_expr + compiled = [] + + def wrapped_set_pref_expr(self, expression, inputs, *args, **kwargs): + res = original_set_pref_expr(self, expression, inputs, *args, **kwargs) + compiled.append(expression) + return res + + monkeypatch.setattr(blosc2.NDArray, "_set_pref_expr", wrapped_set_pref_expr) + return compiled + + @pytest.mark.parametrize( ("dtype", "x", "y"), [ @@ -2450,19 +2469,7 @@ def test_cumulative_block_layout(monkeypatch, blocks): ) @pytest.mark.parametrize("getitem", [False, True]) def test_where_scalar_uses_miniexpr(monkeypatch, dtype, x, y, getitem): - import importlib - - lazyexpr_mod = importlib.import_module("blosc2.lazyexpr") - monkeypatch.setattr(lazyexpr_mod, "try_miniexpr", True) - original_set_pref_expr = blosc2.NDArray._set_pref_expr - compiled = [] - - def wrapped_set_pref_expr(self, expression, inputs, *args, **kwargs): - res = original_set_pref_expr(self, expression, inputs, *args, **kwargs) - compiled.append(expression) - return res - - monkeypatch.setattr(blosc2.NDArray, "_set_pref_expr", wrapped_set_pref_expr) + compiled = _record_miniexpr(monkeypatch) na = np.arange(40 * 30, dtype=dtype).reshape(40, 30) nb = (na // 3).astype(dtype) a = blosc2.asarray(na, chunks=(16, 30), blocks=(4, 30)) @@ -2476,3 +2483,171 @@ def wrapped_set_pref_expr(self, expression, inputs, *args, **kwargs): # The scalar is inlined, so miniexpr (rather than the numexpr fallback) did the work assert len(compiled) == 1 assert "_where_x" in compiled[0] or "_where_y" in compiled[0] + + +@pytest.mark.parametrize( + "case", ["mask", "mask_and_threshold", "labels", "labels_arrays", "wider_output", "int16"] +) +@pytest.mark.parametrize("getitem", [False, True]) +def test_where_mixed_dtypes_use_miniexpr(monkeypatch, case, getitem): + # Conditions whose operands cast safely to the output dtype stay on miniexpr + compiled = _record_miniexpr(monkeypatch) + rng = np.random.default_rng(1) + na = rng.normal(100, 50, (40, 30)).astype(np.float32) + nb = rng.normal(size=(40, 30)).astype(np.float32) + nm = na > 120 + nl = rng.integers(0, 5, (40, 30)).astype(np.uint8) + ni = rng.integers(-100, 100, (40, 30)).astype(np.int16) + a, b, m, lbl, i = (blosc2.asarray(v, chunks=(16, 30), blocks=(4, 30)) for v in (na, nb, nm, nl, ni)) + cond, x, y, ncond, nx, ny = { + "mask": (m, a, 0, nm, na, 0), + "mask_and_threshold": (m & (a > 100), a, 0, nm & (na > 100), na, 0), + "labels": (lbl == 3, a, 0, nl == 3, na, 0), + "labels_arrays": (lbl == 3, a, b, nl == 3, na, nb), + "wider_output": (a > 100, 1.0, 0.0, na > 100, 1.0, 0.0), + "int16": (i > 5, a, i, ni > 5, na, ni), + }[case] + expr = blosc2.where(cond, x, y) + expected = np.where(ncond, nx, ny) + res = expr[:] if getitem else expr.compute()[:] + assert res.dtype == expected.dtype + np.testing.assert_array_equal(res, expected) + assert len(compiled) == 1 + assert "where(" in compiled[0] + + +def test_where_float32_condition_inexact_literal(): + # NumPy compares float32 against float32(0.1), so nothing is > 0.1 here. A fused where() must + # not compare against the double 0.1 (or against 0, the literal typed by an int64 output). + a = blosc2.asarray(np.full(12, 0.1, dtype=np.float32), chunks=(6,), blocks=(3,)) + b = blosc2.asarray(np.full(12, 0.5, dtype=np.float32), chunks=(6,), blocks=(3,)) + for x, y in ((1, 0), (a, 0), (a, b), (1.0, b)): + expr = blosc2.where(a > 0.1, x, y) + expected = np.full(12, y if np.isscalar(y) else 0.5) + np.testing.assert_array_equal(expr[:], expected) + np.testing.assert_array_equal(expr.compute()[:], expected) + expr = blosc2.where(a <= 0.1, x, y) + expected = np.full(12, x if np.isscalar(x) else 0.1, dtype=expr.dtype) + np.testing.assert_array_equal(expr[:], expected) + np.testing.assert_array_equal(expr.compute()[:], expected) + + +def test_where_large_uint64_scalar(): + a = blosc2.asarray(np.arange(12, dtype=np.uint64), chunks=(6,), blocks=(3,)) + big = np.uint64(2**63 + 1) + expr = blosc2.where(a > 5, big, a) + expected = np.array(list(range(6)) + [2**63 + 1] * 6, dtype=np.uint64) + for res in (expr[:], expr.compute()[:]): + assert res.dtype == np.uint64 + np.testing.assert_array_equal(res, expected) + + +@pytest.mark.parametrize("dtype", [np.int64, np.uint64]) +@pytest.mark.parametrize( + "case", ["eq_literal", "add_literal", "add_scalar", "where_scalar", "where_literal_condition"] +) +def test_large_integer_constants_exact(dtype, case): + # miniexpr parses literals into C doubles, so integers beyond 2**53 must not reach it + na = (2**53 + np.arange(-6, 6)).astype(dtype) + a = blosc2.asarray(na, chunks=(6,), blocks=(3,)) + big = 2**53 + 1 + huge = dtype(2**63 + 1 if dtype == np.uint64 else 2**62 + 1) + expr, expected = { + "eq_literal": (a == big, na == big), + "add_literal": (a + big, na + big), + "add_scalar": (a + huge, na + huge), + "where_scalar": (blosc2.where(a > 2**53, huge, a), np.where(na > 2**53, huge, na)), + "where_literal_condition": (blosc2.where(a == big, a, 0), np.where(na == big, na, 0)), + }[case] + for res in (expr[:], expr.compute()[:]): + assert res.dtype == expected.dtype + np.testing.assert_array_equal(res, expected) + + +def _values_around(literal, dtype): + """12 values of `dtype` at and next to `literal`, where comparing with it gets decided.""" + dtype = np.dtype(dtype) + if dtype.kind == "f": + v = dtype.type(literal) + values = [v, np.nextafter(v, dtype.type(np.inf)), np.nextafter(v, dtype.type(-np.inf)), 0] + else: + info = np.iinfo(dtype) + lo, hi = math.floor(literal), math.ceil(literal) + values = [v for v in (lo - 1, lo, hi, hi + 1, 0) if info.min <= v <= info.max] + return np.resize(np.array(values, dtype=dtype), 12) + + +def _where_branches(kind, out_dtype): + """NumPy branches x, y for where(): `out_dtype` scalars, arrays, or an array and a Python 0.""" + if kind == "python": + return (1, 0) if out_dtype == "int" else (0.5, 0.0) + dtype = np.dtype(out_dtype) + # 64-bit integers next to 2**53, to check that their values do not go through doubles + start = 2**53 if dtype.kind in "iu" and dtype.itemsize == 8 else 1 + x = (np.arange(12) % 2).astype(dtype) if dtype.kind == "b" else (start + np.arange(12)).astype(dtype) + y = np.zeros(12, dtype=dtype) + return {"scalars": (x[1], y[0]), "arrays": (x, y), "mixed": (x, 0)}[kind] + + +@pytest.mark.skipif(np.__version__.startswith("1."), reason="NumPy < 2.0 has different casting rules") +@pytest.mark.parametrize( + "cond_dtype", [np.float32, np.float64, np.int8, np.int32, np.int64, np.uint8, np.uint64] +) +@pytest.mark.parametrize("branches", ["python", "scalars", "arrays", "mixed"]) +def test_where_literal_condition_matches_numpy(cond_dtype, branches): + # A fused where() gets typed and computed in its output dtype by miniexpr; comparing against + # literals that the condition's dtype (or a double) cannot hold must still give NumPy's results. + literals = [0.1, 0.5, -1, 100, 16777217, 2**31 + 1, 2**53 + 1] + out_dtypes = ( + ["int", "float"] if branches == "python" else ["float32", "float64", "int64", "uint8", "bool"] + ) + failures = [] + for literal, op, out_dtype in itertools.product(literals, (operator.gt, operator.eq), out_dtypes): + nc = _values_around(literal, cond_dtype) + c = blosc2.asarray(nc, chunks=(6,), blocks=(3,)) + nx, ny = _where_branches(branches, out_dtype) + x, y = ( + blosc2.asarray(v, chunks=(6,), blocks=(3,)) if isinstance(v, np.ndarray) else v for v in (nx, ny) + ) + expected = np.where(op(nc, literal), nx, ny) + expr = blosc2.where(op(c, literal), x, y) + for how, res in (("[:]", expr[:]), ("compute()", expr.compute()[:])): + if res.dtype != expected.dtype or not np.array_equal(res, expected): + failures.append(f"{how} {op.__name__}({literal}) -> {out_dtype}: {res} != {expected}") + assert not failures, "\n".join(failures[:10]) + + +@pytest.mark.skipif(np.__version__.startswith("1."), reason="NumPy < 2.0 has different casting rules") +@pytest.mark.parametrize("cond", ["float32_sum", "int8_wraparound"]) +@pytest.mark.parametrize("branches", ["python_int", "python_float", "array_scalar", "arrays"]) +def test_where_arithmetic_condition_matches_numpy(cond, branches): + # miniexpr would compute the condition in the output dtype: float32 sums in float64, int8 sums + # without wrapping around + rng = np.random.default_rng(0) + if cond == "float32_sum": + na = rng.random(24, dtype=np.float32) + nb = rng.random(24, dtype=np.float32) + operands = (na, nb, na + nb) + + def build(a, b, c): + return a + b == c + else: + na = rng.integers(-128, 128, 24).astype(np.int8) + operands = (na,) + + def build(a): + return a + a < 0 + + to_blosc2 = {id(v): blosc2.asarray(v, chunks=(12,), blocks=(4,)) for v in operands} + nx, ny = { + "python_int": (1, 0), + "python_float": (1.0, 0.0), + "array_scalar": (na, 0), + "arrays": (na, operands[-1]), + }[branches] + x, y = (to_blosc2.get(id(v), v) for v in (nx, ny)) + expected = np.where(build(*operands), nx, ny) + expr = blosc2.where(build(*(to_blosc2[id(v)] for v in operands)), x, y) + for res in (expr[:], expr.compute()[:]): + assert res.dtype == expected.dtype + np.testing.assert_array_equal(res, expected) From 5f6732bca72f71e06df925218c85797656154f66 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Tue, 3 Nov 2026 16:10:49 +0100 Subject: [PATCH 3/3] Keep float32 std/var with ddof float32 on NumPy 1.x; int64 test inputs Two CI failures of the previous commits: - NumPy 1.26: a.std(ddof=1) and a.var(ddof=1) of float32 returned float64. The full std/var returns a float32 scalar, and under NumPy 1.x's value-based casting a float32 scalar times the Python int count of the n / (n - ddof) rescaling becomes float64 (with keepdims, a float32 array times a count above 65535 does too). The rescaled variance is now cast back to its own dtype; with NumPy 2 this changes nothing. This fixes the 24 test_std_var_full_reduction[...-1-...- float32] failures. - WebAssembly: test_large_integer_constants_exact and test_where_literal_condition_matches_numpy added 2**53 to a plain np.arange(), which is int32 on wasm32 and raised OverflowError. They build int64 inputs explicitly now. Not changed here: on NumPy 1.26, mean() of float32 already returns float64 for full reductions and for axis reductions over more than 65535 elements, and mean(where=) does on NumPy 2 too (also on 7db7b09). Co-Authored-By: Claude Opus 5.5 --- src/blosc2/lazyexpr.py | 4 +++- tests/ndarray/test_lazyexpr.py | 6 ++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 8e40e9e9b..1487a6119 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -4538,7 +4538,9 @@ def _var(self, axis, dtype, keepdims, ddof, where, fp_accuracy, kwargs, finish=N axis=axis, dtype=np.int64, fp_accuracy=fp_accuracy, **execution_kwargs ) ) - out = out * num_elements / (num_elements - ddof) + # Keep the dtype of the variance: with NumPy 1.x, a float32 scalar times a Python int + # becomes float64, and so does a float32 array times an int above 65535 + out = (out * num_elements / (num_elements - ddof)).astype(np.asarray(out).dtype, copy=False) if finish is not None: out = finish(out) out2 = kwargs.pop("out", None) diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index c2e38bb3c..6837a7b46 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -2548,7 +2548,7 @@ def test_where_large_uint64_scalar(): ) def test_large_integer_constants_exact(dtype, case): # miniexpr parses literals into C doubles, so integers beyond 2**53 must not reach it - na = (2**53 + np.arange(-6, 6)).astype(dtype) + na = (2**53 + np.arange(-6, 6, dtype=np.int64)).astype(dtype) # the default int is int32 on wasm32 a = blosc2.asarray(na, chunks=(6,), blocks=(3,)) big = 2**53 + 1 huge = dtype(2**63 + 1 if dtype == np.uint64 else 2**62 + 1) @@ -2584,7 +2584,9 @@ def _where_branches(kind, out_dtype): dtype = np.dtype(out_dtype) # 64-bit integers next to 2**53, to check that their values do not go through doubles start = 2**53 if dtype.kind in "iu" and dtype.itemsize == 8 else 1 - x = (np.arange(12) % 2).astype(dtype) if dtype.kind == "b" else (start + np.arange(12)).astype(dtype) + # int64 explicitly: the default integer is int32 on wasm32, where 2**53 would overflow it + values = np.arange(12, dtype=np.int64) + x = (values % 2).astype(dtype) if dtype.kind == "b" else (start + values).astype(dtype) y = np.zeros(12, dtype=dtype) return {"scalars": (x[1], y[0]), "arrays": (x, y), "mixed": (x, 0)}[kind]