Align DSL operand grids, so mixed-dtype NDArray inputs work

Blocks are sized in bytes, so same-shaped operands of different itemsize
get different chunks/blocks (float32 vs int64 over 1M elements: blocks of
31250 vs 15625). validate_inputs then refuses the miniexpr fast path, and
a DSL kernel has no slow path to fall back on, so evaluation died with a
misleading "slicing a DSL computation is not supported".

Whether the grids diverge depends on array size and on the platform's
cache detection, which is why CI caught it only on Windows (where
platform.machine() is "AMD64", so compute_chunks_blocks' x86_64 branch
never runs and blocks stay L1-sized) and WASM: their 10,007-element
operands already disagree, while elsewhere both fit a single block. At 1M
elements it failed everywhere, macOS included -- hence the regression test
sits at that size.

LazyUDF now copies mismatched NDArray operands onto the grid of the widest
dtype (fewest elements per block, so every operand still fits the cache
budget the heuristic aimed at), once at construction rather than per
evaluation, and only for DSL kernels.

Also trace the JS bridge under BLOSC_ME_JIT_TRACE. It bypasses miniexpr
entirely, so it used to report nothing at all, which is what turned the
WASM-only mandel dispatch test into a confusing empty-output failure
rather than an obvious one. Both engines now trace, and that test asserts
on both platforms instead of skipping.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: FrancescAlted

src/blosc2/lazyexpr.py

@@ -1495,6 +1495,15 @@ def _js_dtypes_ok(operands, kwargs) -> bool:
     )
 
 
+def _trace_js_backend(expression):
+    """BLOSC_ME_JIT_TRACE counterpart for the JS bridge, which never reaches
+    miniexpr's trace point in `fast_eval` (see there for the message format)."""
+    if os.environ.get("BLOSC_ME_JIT_TRACE", "").lower() in ("1", "true", "on"):
+        source = getattr(expression, "dsl_source", None) or expression
+        expr_short = str(source)[:120].replace("\n", " ")
+        print(f"[blosc2] engine=js expr={expr_short}", flush=True)
+
+
 def _maybe_js_backend(expression, jit, jit_backend, reduce_args, operands, kwargs, shape=None):
     """Resolve the JS backend for a DSL kernel.
 
@@ -1524,7 +1533,9 @@ def _maybe_js_backend(expression, jit, jit_backend, reduce_args, operands, kwarg
                 'jit_backend="js" requires a floating-point output dtype '
                 f"(got {np.dtype(out_dtype)}); drop jit_backend to use miniexpr"
             )
-        return _as_js_udf(expression, shape), None, None
+        bridge = _as_js_udf(expression, shape)
+        _trace_js_backend(expression)
+        return bridge, None, None
     prefer_js = (
         jit is not False  # jit=True/None prefer the best JIT (js); only jit=False forces interpreter
         and jit_backend is None
@@ -1541,6 +1552,7 @@ def _maybe_js_backend(expression, jit, jit_backend, reduce_args, operands, kwarg
         bridge = _as_js_udf(expression, shape)  # transpiles; raises on any unsupported construct
     except Exception:
         return expression, jit, jit_backend  # fall back to miniexpr, no regression
+    _trace_js_backend(expression)
     return bridge, None, None
 
 
@@ -4576,12 +4588,45 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No
         return new_expr
 
 
+def _align_dsl_operand_grids(inputs):
+    """Put every NDArray operand of a DSL kernel on a single chunks/blocks grid.
+
+    Blocks are sized in bytes, so same-shaped operands of different itemsize get
+    different grids by default (float32 vs int64 over 1M elements: blocks of
+    31250 vs 15625 elements).  miniexpr needs one common grid, and a DSL kernel
+    has no slow path to fall back on, so evaluation would fail outright with a
+    confusing "slicing is not supported" error.  Copy the odd operands onto the
+    grid of the widest dtype: its blocks hold the fewest elements, so every
+    operand still fits within the cache budget the heuristic aimed at.
+
+    The copies happen once, at construction, and only when the grids actually
+    disagree -- whether they do depends on the array size and on the platform's
+    cache detection, which is why this used to fail only on some CI runners.
+    """
+    nd = [x for x in inputs if isinstance(x, blosc2.NDArray) and x.ndim > 0]
+    if len(nd) < 2 or len({(x.shape, x.chunks, x.blocks) for x in nd}) < 2:
+        return inputs
+    ref = max(nd, key=lambda x: x.dtype.itemsize)
+    aligned = []
+    for x in inputs:
+        misaligned = (
+            isinstance(x, blosc2.NDArray)
+            and x.ndim > 0
+            and x.shape == ref.shape
+            and (x.chunks, x.blocks) != (ref.chunks, ref.blocks)
+        )
+        aligned.append(x.copy(chunks=ref.chunks, blocks=ref.blocks) if misaligned else x)
+    return aligned
+
+
 class LazyUDF(LazyArray):
     def __init__(
         self, func, inputs, dtype, shape=None, chunked_eval=True, jit=None, jit_backend=None, **kwargs
     ):
         # After this, all the inputs should be np.ndarray or NDArray objects
         self.inputs = convert_inputs(inputs)
+        if isinstance(func, DSLKernel):
+            self.inputs = _align_dsl_operand_grids(self.inputs)
         # Get res shape
         if shape is None:
             self._shape = compute_broadcast_shape(self.inputs)

tests/ndarray/test_dsl_kernels.py

@@ -1196,6 +1196,20 @@ def test_dsl_kernel_numpy_operands_mixed_dtype_promotes_output():
     np.testing.assert_array_equal(res, ref)
 
 
+def test_dsl_kernel_ndarray_operands_with_different_itemsize():
+    # Blocks are sized in bytes, so a float32 and an int64 operand get different
+    # chunks/blocks by default; the DSL path has no slow fallback, so it used to
+    # raise "slicing is not supported" whenever the grids diverged (which depends
+    # on array size and on the platform's cache detection).
+    n = 1_000_000
+    a = (np.arange(n) % 7).astype(np.float32)
+    b = (np.arange(n) % 5).astype(np.int64)
+    A, B = blosc2.asarray(a), blosc2.asarray(b)
+    assert (A.chunks, A.blocks) != (B.chunks, B.blocks)
+    res = blosc2.lazyudf(_numpy_operand_kernel, (A, B), dtype=None)[()]
+    np.testing.assert_array_equal(res, a * 2.0 + b)
+
+
 def test_dsl_kernel_mixed_ndarray_and_numpy_operand():
     shape = (20, 10)
     a = np.arange(np.prod(shape), dtype=np.float64).reshape(shape)

tests/ndarray/test_jit_dsl_dispatch.py

@@ -53,7 +53,9 @@ def mandel(cr, ci, max_iter):
     monkeypatch.setenv("BLOSC_ME_JIT_TRACE", "1")
     res = mandel(cr, ci, 30)
     captured = capsys.readouterr()
-    assert "engine=miniexpr" in captured.out
+    # Under WebAssembly this kernel is transpiled to the JS bridge instead of
+    # going through miniexpr (see _maybe_js_backend); both engines trace.
+    assert f"engine={'js' if blosc2.IS_WASM else 'miniexpr'}" in captured.out
     assert "def mandel" in captured.out
     np.testing.assert_array_equal(res, _mandel_numpy(cr, ci, 30))