Ignore dims with stride zero in elementwise pipeline, fix docs typos

README.md

@@ -14,7 +14,7 @@ Only the blocks that overlap a read request are decompressed, so random access i
 
 **Lazy operation chains.** Every operation - arithmetic, shape change, reduction, type cast - returns a new view that wraps the
 input(s) and records the transformation, nothing is computed until data is explicitly requested.
-A chain of such operations build a pipeline that runs in a single decompression pass the moment you ask for output.
+A chain of such operations builds a pipeline that runs in a single decompression pass the moment you ask for output.
 
 ```rust
 use jix::Array;
@@ -83,7 +83,7 @@ can be used independently, and each fits a different scenario.
 
     In Rust, plain iterators over regular ndarrays already
     give you lazy element-wise evaluation for free (although naive use of the `ndarray` crate may still
-    produces NumPy-style intermediates), so for simple `map`/`zip`-style pipelines
+    produce NumPy-style intermediates), so for simple `map`/`zip`-style pipelines
     jix offers little over hand-written iterator code.
     The advantage shows up once the pipeline includes
     operations that change the shape or the access pattern - reductions, broadcasts,

jix-py/README.md

@@ -7,7 +7,7 @@ Only the blocks that overlap a read request are decompressed, so random access i
 
 **Lazy operation chains.** Every operation - arithmetic, shape change, reduction, type cast - returns a new view that wraps the
 input(s) and records the transformation, nothing is computed until data is explicitly requested.
-A chain of such operations build a pipeline that runs in a single decompression pass the moment you ask for output.
+A chain of such operations builds a pipeline that runs in a single decompression pass the moment you ask for output.
 
 ```python
 import jix

jix-py/docs/index.md

@@ -7,7 +7,7 @@ Only the blocks that overlap a read request are decompressed, so random access i
 
 **Lazy operation chains.** Every operation - arithmetic, shape change, reduction, type cast - returns a new view that wraps the
 input(s) and records the transformation, nothing is computed until data is explicitly requested.
-A chain of such operations build a pipeline that runs in a single decompression pass the moment you ask for output.
+A chain of such operations builds a pipeline that runs in a single decompression pass the moment you ask for output.
 
 ```python
 import jix

jix-py/docs/module.md

@@ -10,7 +10,7 @@ Only the blocks that overlap a read request are decompressed, so random access i
 
 **Lazy operation chains.** Every operation - arithmetic, shape change, reduction, type cast - returns a new view that wraps the
 input(s) and records the transformation, nothing is computed until data is explicitly requested.
-A chain of such operations build a pipeline that runs in a single decompression pass the moment you ask for output.
+A chain of such operations builds a pipeline that runs in a single decompression pass the moment you ask for output.
 
 The library is NumPy-compatible: arrays expose a NumPy `dtype`, accept NumPy index syntax,
 and materialize to NumPy arrays on demand.
@@ -86,7 +86,7 @@ output.
 | Function | Description |
 |---|---|
 | `jix.compact(...)` | Compress any array-like (NumPy array, list, scalar) into a new jix array. This is the primary constructor. |
-| `jix.asarray(...)` | Wrap any array-like as a zero-copy jix view without compressing. Useful for mixing plain NumPy data with jix arrays in ations. |
+| `jix.asarray(...)` | Wrap any array-like as a zero-copy jix view without compressing. Useful for mixing plain NumPy data with jix arrays in operations. |
 | `jix.read_array(...)` | Load a `.jix` file from disk. |
 
 **Reading data from an `Array`:**
@@ -172,7 +172,7 @@ b = jix.read_array("data.jix")
 c = jix.read_array("data.jix", mmap=True)
 ```
 
-`write_array` accept a file path or any writable binary
+`write_array` accepts a file path or any writable binary
 file-like object. `read_array` accepts a file path or any seekable binary file-like
 object.
 
@@ -185,5 +185,5 @@ disk.
 # Limits
 
 - Maximum array dimensions: 8.
-- Maximum inner-shape dimensions for struct dtypes: 4.
+- Maximum inner-shape dimensions for dtypes: 4.
 - Little-endian platforms only.

jix-py/python/tests/test_broadcasting.py

@@ -269,3 +269,158 @@ def test_broadcast_sub_index(shape_a, shape_b, idx):
     expected = np_a + np_b
 
     np.testing.assert_array_equal(result[idx], expected[idx])
+
+
+# ---------------------------------------------------------------------------
+# where
+# ---------------------------------------------------------------------------
+
+
+@st.composite
+def broadcastable_shapes_triple(draw, max_ndim: int = 4, max_dim: int = 5):
+    """
+    Generate three shapes that are mutually broadcastable per numpy rules.
+
+    Same right-aligned construction as broadcastable_shapes_pair, extended to
+    three operands: at each position every operand that has the dimension gets
+    either the common size or 1.
+    """
+    ndims = [draw(st.integers(1, max_ndim)) for _ in range(3)]
+    ndim = max(ndims)
+    dims = [[] for _ in range(3)]
+
+    for pos in range(ndim):
+        d = draw(st.integers(1, max_dim))
+        for i, nd in enumerate(ndims):
+            if pos >= ndim - nd:
+                dims[i].append(d if draw(st.booleans()) else 1)
+
+    return tuple(tuple(d) for d in dims)
+
+
+_WHERE_SHAPE_TRIPLES = [
+    # all equal
+    ((3,), (3,), (3,)),
+    # one operand is a length-1 axis
+    ((3,), (3,), (1,)),
+    ((1,), (3,), (3,)),
+    ((3,), (1,), (3,)),
+    # condition drives a leading axis, y is a row vector
+    ((2, 1), (2, 3), (3,)),
+    ((2, 3), (1, 3), (2, 1)),
+    # fewer dims get prepended
+    ((4,), (3, 4), (3, 1)),
+    ((3, 1, 4), (3, 5, 1), (5, 4)),
+    # 0-d operands broadcast against everything
+    ((), (2, 3), (2, 3)),
+    ((2, 3), (), (2, 3)),
+    ((2, 3), (2, 3), ()),
+    # classic numpy docs shapes
+    ((8, 1, 6, 1), (7, 1, 5), (6, 5)),
+]
+
+
+@pytest.mark.parametrize("cond_shape, x_shape, y_shape", _WHERE_SHAPE_TRIPLES)
+def test_where_broadcast_shapes(cond_shape, x_shape, y_shape):
+    """jix.where broadcasts condition/x/y exactly like numpy.where."""
+    np_cond = (np.arange(np.prod(cond_shape, dtype=int)) % 2 == 0).reshape(cond_shape)
+    np_x = np.arange(1, np.prod(x_shape, dtype=int) + 1, dtype=np.int32).reshape(x_shape)
+    np_y = -np.arange(1, np.prod(y_shape, dtype=int) + 1, dtype=np.int32).reshape(y_shape)
+
+    result = jix.where(jix.compact(np_cond), jix.compact(np_x), jix.compact(np_y))
+    expected = np.where(np_cond, np_x, np_y)
+
+    assert result.shape == expected.shape, f"shape: {result.shape} != {expected.shape}"
+    assert result.dtype == np_x.dtype
+    np.testing.assert_array_equal(result.numpy(), expected)
+
+
+@given(st.data())
+def test_where_broadcast_property(data: DataObject):
+    """jix.where matches numpy.where for arbitrary mutually broadcastable shapes."""
+    cond_shape, x_shape, y_shape = data.draw(broadcastable_shapes_triple(), label="shapes")
+
+    np_cond = (np.arange(np.prod(cond_shape, dtype=int)) % 3 != 0).reshape(cond_shape)
+    np_x = np.arange(1, np.prod(x_shape, dtype=int) + 1, dtype=np.int32).reshape(x_shape)
+    np_y = -np.arange(1, np.prod(y_shape, dtype=int) + 1, dtype=np.int32).reshape(y_shape)
+
+    result = jix.where(jix.compact(np_cond), jix.compact(np_x), jix.compact(np_y))
+    expected = np.where(np_cond, np_x, np_y)
+
+    assert result.shape == expected.shape
+    np.testing.assert_array_equal(result.numpy(), expected)
+
+
+@pytest.mark.parametrize(
+    "cond_shape, x_shape, y_shape",
+    [
+        ((3,), (3,), (4,)),
+        ((3,), (4,), (3,)),
+        ((4,), (3,), (3,)),
+        ((2, 3), (2, 3), (3, 2)),
+    ],
+)
+def test_where_broadcast_incompatible_raises(cond_shape, x_shape, y_shape):
+    """Shapes that numpy cannot broadcast together raise an error."""
+    cond = jix.compact(np.ones(cond_shape, dtype=bool))
+    x = jix.compact(np.ones(x_shape, dtype=np.int32))
+    y = jix.compact(np.ones(y_shape, dtype=np.int32))
+
+    with pytest.raises(Exception):
+        _ = jix.where(cond, x, y).numpy()
+
+
+def test_where_broadcast_numpy_scalar_operands():
+    """numpy scalar x/y operands broadcast to the condition shape."""
+    np_cond = np.array([True, False, True])
+    result = jix.where(jix.compact(np_cond), np.int32(1), np.int32(0))
+    expected = np.where(np_cond, np.int32(1), np.int32(0))
+
+    assert result.shape == expected.shape
+    np.testing.assert_array_equal(result.numpy(), expected)
+
+
+def test_where_broadcast_python_scalar_operands():
+    """Python int x/y operands broadcast to the condition shape."""
+    np_cond = np.array([[True, False], [False, True]])
+    result = jix.where(jix.compact(np_cond), 1, 0)
+
+    np.testing.assert_array_equal(result.numpy(), np.where(np_cond, 1, 0))
+
+
+def test_where_broadcast_scalar_condition():
+    """A 0-d condition broadcasts against the value arrays."""
+    np_x = np.array([1, 2, 3], dtype=np.int32)
+    np_y = np.array([10, 20, 30], dtype=np.int32)
+
+    for cond in (True, False):
+        result = jix.where(jix.compact(np.array(cond)), jix.compact(np_x), jix.compact(np_y))
+        np.testing.assert_array_equal(result.numpy(), np.where(cond, np_x, np_y))
+
+
+def test_where_broadcast_lazy_operands():
+    """Broadcasting composes with lazy views on any of the three operands."""
+    np_x = np.arange(1, 7, dtype=np.int32).reshape(2, 3)
+    np_y = np.array([[100], [200]], dtype=np.int32)
+    x = jix.compact(np_x)
+    y = jix.compact(np_y)
+    cond = jix.compact(np.array([True, False, True]))
+
+    result = jix.where(cond, x * 2, y)
+    expected = np.where(np.array([True, False, True]), np_x * 2, np_y)
+
+    assert result.shape == expected.shape
+    np.testing.assert_array_equal(result.numpy(), expected)
+
+
+def test_where_broadcast_sub_index():
+    """Slicing a broadcast where-result gives the same values as numpy."""
+    np_cond = np.array([[True], [False], [True]])
+    np_x = np.arange(1, 13, dtype=np.int32).reshape(3, 4)
+    np_y = np.array([-1, -2, -3, -4], dtype=np.int32)
+
+    result = jix.where(jix.compact(np_cond), jix.compact(np_x), jix.compact(np_y))
+    expected = np.where(np_cond, np_x, np_y)
+
+    idx = (slice(1, 3), slice(1, 4))
+    np.testing.assert_array_equal(result[idx], expected[idx])

jix-py/python/tests/test_type_promotion.py

@@ -498,3 +498,64 @@ def test_power_rejects_complex():
     b = jix.compact(np.array([1 + 0j, 2 + 0j], dtype=np.complex64))
     with pytest.raises(Exception):
         _ = jix.power(a, b).numpy()
+
+
+# ---------------------------------------------------------------------------
+# Scalar operands are built directly at the dispatch dtype.
+#
+# Dispatch casts every operand to the dtype of the impl it picked. For an array
+# that means wrapping it in a lazy `Cast` view, but a scalar has nothing to wrap:
+# the value is cast up front and the `Scalar` storage is created at the target
+# dtype. A `Cast<Scalar>` in the storage chain means that stopped happening.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("arr_dtype", [np.int8, np.int32, np.int64, np.float32, np.float64])
+@pytest.mark.parametrize("scalar", [2, 2.0, True, np.float32(2.0), np.int8(2), np.float16(2.0)])
+def test_scalar_operand_has_no_cast_view(arr_dtype, scalar):
+    a = jix.compact(np.array([1, 2, 3], dtype=arr_dtype))
+    storage = repr(a * scalar)
+    assert "Scalar" in storage, storage
+    # `Cast<Compact>` is fine - the array operand really does need a lazy cast view.
+    assert "Cast<Scalar" not in storage, storage
+
+
+def test_scalar_operand_has_no_cast_view_for_cmp_and_bitwise():
+    a = jix.compact(np.array([1, 2, 3], dtype=np.float32))
+    assert "Cast<Scalar" not in repr(a > 2.0)
+    b = jix.compact(np.array([1, 2, 3], dtype=np.int32))
+    assert "Cast<Scalar" not in repr(b | 1)
+
+
+def test_scalar_operand_complex_has_no_cast_view():
+    a = jix.compact(np.array([1 + 2j, 3 + 4j], dtype=np.complex64))
+    storage = repr(a * 2.0)
+    assert "Scalar" in storage, storage
+    assert "Cast<Scalar" not in storage, storage
+
+
+def test_python_float_scalar_rounds_to_dispatch_dtype():
+    """`0.1` is not representable in f32: the scalar must round exactly the way numpy's
+    float32(0.1) does, whether it is cast before or after the Scalar storage is built."""
+    arr = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    a = jix.compact(arr)
+    np.testing.assert_array_equal((a * 0.1).numpy(), arr * np.float32(0.1))
+
+
+def test_float16_scalar_keeps_its_value():
+    arr = np.array([1.0, 2.0, 3.0], dtype=np.float16)
+    a = jix.compact(arr)
+    np.testing.assert_array_equal((a * np.float16(0.1)).numpy(), arr * np.float16(0.1))
+
+
+def test_bool_scalar_keeps_its_value():
+    arr = np.array([1, 2, 3], dtype=np.int32)
+    a = jix.compact(arr)
+    np.testing.assert_array_equal((a * True).numpy(), arr * np.int32(1))
+    np.testing.assert_array_equal((a * False).numpy(), arr * np.int32(0))
+
+
+def test_complex_scalar_keeps_its_value():
+    arr = np.array([1 + 2j, 3 + 4j], dtype=np.complex64)
+    a = jix.compact(arr)
+    np.testing.assert_array_equal((a * (0.1 + 0.2j)).numpy(), arr * np.complex64(0.1 + 0.2j))

jix-py/src/array.rs

@@ -329,7 +329,7 @@ impl Array {
     /// The total number of elements in the array (the product of the axis lengths).
     ///
     /// Returns:
-    ///     The total element n_printed as an integer.
+    ///     The total element count as an integer.
     ///
     /// ```python
     /// import jix
@@ -1293,7 +1293,7 @@ impl Array {
         crate::ops::clamp(slf, min, max)
     }
 
-    /// Returns the sign of each element as a floating-point value. See [`jix.sign()`][jix.sign].
+    /// Returns the sign of each element. See [`jix.sign()`][jix.sign].
     pub fn sign(slf: &Bound<'_, Self>) -> PyResult<Self> {
         crate::ops::sign(slf)
     }

jix-py/src/lib.rs

@@ -3,7 +3,7 @@
 //
 #![doc = include_str!("../docs/module.md")]
 //!
-//! # Disclaimer
+//! # Acknowledgements
 //!
 //! This project would not exist without the work of several upstream authors and communities.
 //! Specifically, this project was greatly inspired by the [C-Blosc2](https://github.com/Blosc/c-blosc2) library.

jix-py/src/ops/as_array.rs

@@ -10,8 +10,8 @@ use crate::ops::common::Operand;
 
 /// Convert any array-like object to a [`jix.Array`][jix.Array].
 ///
-/// This function differ from [`jix.compact()`][jix.compact] in that it does not compress the data -
-/// it always produces a view of the input data. In Some cases a copy may be necessary, for example
+/// This function differs from [`jix.compact()`][jix.compact] in that it does not compress the data -
+/// it always produces a view of the input data. In some cases a copy may be necessary, for example
 /// to convert from a raw python list to a typed buffer, but in general this function tries to avoid
 /// copying data when possible, and it never compresses the data.
 ///

jix-py/src/ops/bitwise.rs

@@ -318,7 +318,8 @@ define_op2!(
     ///
     /// Shifts the bits of each element of `a` left by the corresponding value in `b`.
     /// Vacated bits are filled with zeros. Shifting by a value greater than or equal to the
-    /// bit width of the type produces zero.
+    /// bit width of the type panics in debug builds and masks the shift amount modulo the bit
+    /// width in release builds (it does NOT produce zero).
     ///
     /// **Type promotion**: if `a` and `b` have different integer dtypes, both are cast to
     /// the smallest integer type that can represent both (Safe casting rules).
@@ -360,7 +361,8 @@ define_op2!(
     /// For **unsigned** types this is a logical shift: vacated bits are filled with zeros.
     /// For **signed** types this is an arithmetic shift: vacated bits are filled with the
     /// sign bit (the result preserves the sign). Shifting by a value greater than or equal
-    /// to the bit width produces zero (unsigned) or the sign-extended value (signed).
+    /// to the bit width panics in debug builds and masks the shift amount modulo the bit width
+    /// in release builds (it does NOT produce zero).
     ///
     /// **Type promotion**: if `a` and `b` have different integer dtypes, both are cast to
     /// the smallest integer type that can represent both (Safe casting rules).

jix-py/src/ops/common/dispatch.rs

@@ -1,12 +1,11 @@
 use std::borrow::Cow;
 
-use jix_core::dtype::{Dtype, Dtyped, ScalarKind};
+use jix_core::dtype::{Dtyped, ScalarKind};
 use jix_core::ops::IntoType;
 use jix_core::storage::ArrayStorageAny;
 use jix_core::{Array as CoreArray, ArrayAny, Ty};
 use pyo3::prelude::*;
 
-use crate::ops::astype_impl;
 use crate::ops::common::{CastKind, Operand, Precision, Rank, Scalar};
 use crate::util::{IntoPyResult, ItemOrSequence, IterExt};
 
@@ -201,10 +200,7 @@ impl<const IN_N: usize, ExtraArgs> OpDescriptor<IN_N, ExtraArgs> {
                 let inputs: [ArrayAny; IN_N] = inputs
                     .into_iter()
                     .zip(op_fn.input_desc.iter())
-                    .map(|(input, input_desc)| {
-                        let input = input.into_array()?;
-                        astype_impl(input, &Dtype::new_scalar(input_desc.dtype))
-                    })
+                    .map(|(input, input_desc)| input.cast(input_desc.dtype)?.into_array())
                     .try_collect_array()?
                     .unwrap();