From 155e48f27cb054d9f658be44df264ce2713c21c5 Mon Sep 17 00:00:00 2001 From: ajz34 Date: Tue, 6 Oct 2026 18:18:55 +0800 Subject: [PATCH 1/2] rstsr-dtype-traits (fix): restore the missing i8 x i16 promotion impl `impl_promotion_asable!(i8, i16, false, true, i16)` had been swallowed by its section comment since ce977a1 ("Apply type promotion rules", 2025-09-29): the source reads `// internal typeimpl_promotion_asable!(...)`, so `DTypePromoteAPI for i8` (and the `DTypeCastAPI for i8` twin the same macro emits) never existed, while the mirror row `i16 x i8` did. Every `DTypePromoteAPI`-bound op (maximum/minimum, floor_divide, atan2, copysign, hypot, nextafter, logaddexp, comparisons) therefore declined int8 x int16 pairs; the array-API test suite draws both orderings. Restored with the flags of the mirror row (`Res = i16`; NumPy's result_type and the suite's own promotion table agree), and promotion.rs gains an inline test module: a compile-time completeness check over the full 13 x 13 matrix of canonical dtypes, `Res` checks for representative pairs, and value checks for the restored row. Co-authored-by: Claude Code Co-authored-by: deepseek-v4-flash --- rstsr-dtype-traits/src/promotion.rs | 76 ++++++++++++++++++++++++++++- 1 file changed, 75 insertions(+), 1 deletion(-) diff --git a/rstsr-dtype-traits/src/promotion.rs b/rstsr-dtype-traits/src/promotion.rs index 5705bbf0..ce48ec83 100644 --- a/rstsr-dtype-traits/src/promotion.rs +++ b/rstsr-dtype-traits/src/promotion.rs @@ -228,7 +228,8 @@ macro_rules! impl_promotion_asable { }; } -// internal typeimpl_promotion_asable!(i8, i16, false, true, i16); +// internal type +impl_promotion_asable!(i8, i16, false, true, i16); impl_promotion_asable!(i8, i32, false, true, i32); impl_promotion_asable!(i8, i64, false, true, i64); impl_promotion_asable!(i8, u8, false, false, i16); @@ -558,3 +559,76 @@ impl DTypeCastAPI for c32 { } /* #endregion */ + +/* #region tests */ + +#[cfg(test)] +mod tests { + use super::*; + + /// Compile-time completeness check: the pair promotes. + fn promotes() + where + A: DTypePromoteAPI, + { + } + + /// Compile-time check with a known promoted result type. + fn promotes_to() + where + A: DTypePromoteAPI, + { + } + + /// The full ordered matrix over the canonical dtypes must promote. This + /// guards the `i8 x i16` pair (and its `DTypeCastAPI for i8`): on + /// 2025-09-29 a section comment merged with the following impl line and + /// silently dropped it, so every `DTypePromoteAPI`-bound op declined the + /// pair. + #[test] + fn promotion_matrix_is_complete() { + macro_rules! all_pairs { + ($($a:ty => $($b:ty),+ ;)+) => { + $( $( promotes::<$a, $b>(); )+ )+ + }; + } + all_pairs! { + bool => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + i8 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + i16 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + i32 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + i64 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + u8 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + u16 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + u32 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + u64 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + f32 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + f64 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + c32 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + c64 => bool, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, c32, c64; + } + } + + #[test] + fn promotion_results() { + promotes_to::(); + promotes_to::(); + promotes_to::(); + promotes_to::(); + promotes_to::(); + promotes_to::(); + promotes_to::(); + promotes_to::(); + promotes_to::(); + } + + /// Values of the restored `i8 x i16` row (mirror of `i16 x i8`). + #[test] + fn promote_and_cast_i8_i16() { + assert_eq!(>::promote_other(300i16), 300i16); + assert_eq!(>::promote_self(-5i8), -5i16); + assert_eq!(>::into_cast(7i8), 7i16); + } +} + +/* #endregion */ From 8d4cb286b5f6b8e41afa0765a79328ec64aeb53d Mon Sep 17 00:00:00 2001 From: ajz34 Date: Tue, 6 Oct 2026 18:19:03 +0800 Subject: [PATCH 2/2] rstsr-faer-py (feat): bind where over rt::where_f MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The last missing entry of the searching surface (gap G-037): the shim's `where` is a thin wrapper over core `rt::where` (rt::where_f), which landed on main with #114. - Rust: `op_where` calls `where_f` with the output type taken from rstsr's own promotion trait (`TX: DTypePromoteAPI`); the pyfunction is exported as `where` (the Rust keyword forces the `where_` item name; the attribute restores it). - Dispatch: a new `dispatch_where!` — 169 arms, one per dtype pair of the 13 canonical dtypes. Every arm calls the same generic fn item and the result variant follows the promoted type via `any_of`, so the shim carries no promotion table of its own (same construction rule as the W2 binary dispatch macros). - Python: `where(condition, x1, x2, /)` validates a boolean Array condition and marshals x1/x2 through the existing `_operands` weak-scalar path (a scalar becomes a 0-d array of the array operand's dtype, spec "Mixing arrays with Python scalars"); a call with two scalars raises, matching the spec's "at least one of x1 and x2 must be an array". Non-boolean conditions are declined, never truthiness-cast. Suite (array-api-tests @ 6c0b59f): 996 -> 1014 passing, 304 -> 286 failing, 82 skipped — 18 flips / 0 regressions, identical on a canonical FRESH run and a warm re-run. Besides `test_where` and `test_where_with_scalars`, this unblocks the suite's own verification paths that call `xp.where` (getitem, asarray_arrays, eye, tril, triu, linspace, unstack, positive) and five nan-propagation generators. Co-authored-by: Claude Code Co-authored-by: deepseek-v4-flash --- .../rstsr-faer-py/python/rstsr_faer/api.py | 27 +- .../rstsr-faer-py/src/any_tensor.rs | 281 ++++++++++++++++++ crates-interop/rstsr-faer-py/src/lib.rs | 1 + crates-interop/rstsr-faer-py/src/ops.rs | 36 ++- 4 files changed, 342 insertions(+), 3 deletions(-) diff --git a/crates-interop/rstsr-faer-py/python/rstsr_faer/api.py b/crates-interop/rstsr-faer-py/python/rstsr_faer/api.py index 06f9de40..c437199d 100644 --- a/crates-interop/rstsr-faer-py/python/rstsr_faer/api.py +++ b/crates-interop/rstsr-faer-py/python/rstsr_faer/api.py @@ -122,6 +122,7 @@ argmin as _argmin, count_nonzero as _count_nonzero, sum_bool as _sum_bool, + where as _where, take as _take, ) @@ -1317,6 +1318,30 @@ def count_nonzero(x, /, *, axis=None, keepdims=False): return _wrap(_count_nonzero(h, axes, _py_bool(keepdims))) +def where(condition, x1, x2, /): + """Element-wise select: x1 where ``condition`` is True, x2 elsewhere. + + The condition must be an Array with boolean dtype (the spec's "should have + a boolean data type"; non-bool conditions are declined, never truthiness- + cast). ``x1``/``x2`` may each be an Array or a Python scalar; scalars are + converted to 0-d arrays of the array operand's dtype before the native call + (spec "Mixing arrays with Python scalars" — the same weak-scalar rule as + the binary functions, via ``_operands``). At least one of ``x1``/``x2`` + must be an Array; dtype promotion of Array pairs happens rust-side in + rstsr's own matrix. + """ + if not isinstance(condition, Array): + raise TypeError( + f"where: condition must be an rstsr_faer.api Array, got {type(condition).__name__}" + ) + if _kind(condition.dtype) != "bool": + raise TypeError( + f"where: condition must have a boolean data type, got {condition.dtype!r}" + ) + a, b = _operands(x1, x2) + return _wrap(_where(condition._h, a, b)) + + def take(x, /, indices, *, axis=None): h = _handle(x) if not isinstance(indices, Array): @@ -1390,7 +1415,7 @@ def iinfo(type, /): "reshape", "permute_dims", "broadcast_arrays", "broadcast_shapes", "concat", "stack", "unstack", "expand_dims", "squeeze", "flip", "moveaxis", # searching / indexing - "argmax", "argmin", "count_nonzero", "take", + "argmax", "argmin", "count_nonzero", "take", "where", # data types "astype", "finfo", "iinfo", # constants / sentinels diff --git a/crates-interop/rstsr-faer-py/src/any_tensor.rs b/crates-interop/rstsr-faer-py/src/any_tensor.rs index 6e1c6019..ebbef94b 100644 --- a/crates-interop/rstsr-faer-py/src/any_tensor.rs +++ b/crates-interop/rstsr-faer-py/src/any_tensor.rs @@ -986,6 +986,287 @@ macro_rules! dispatch_bin_promote_eq { } pub(crate) use dispatch_bin_promote_eq; +/// Ternary dispatch for the element-wise select (`where`): a boolean +/// condition plus any `x`/`y` dtype pair in rstsr's promotion matrix. Every +/// arm calls the same generic fn item `$f::(cond, a, b)`; the result +/// variant follows the *promoted* type (`any_of`), so no promotion table is +/// duplicated here (arms mirror `rstsr-dtype-traits/src/promotion.rs`, which +/// promotes all 169 pairs of the 13 canonical dtypes). +macro_rules! dispatch_where { + ($cond:expr, $x:expr, $y:expr, $f:ident) => { + match (&$x, &$y) { + (AnyTensor::Bool(a), AnyTensor::Bool(b)) => { + lift(($f::)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::Bool(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::Bool(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::Bool(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::Bool(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::Bool(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::Bool(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::Bool(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::Bool(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::Bool(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::Bool(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::Bool(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::Bool(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::I8(a), AnyTensor::Bool(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I8(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::I8(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::I16(a), AnyTensor::Bool(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I16(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::I16(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::I32(a), AnyTensor::Bool(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I32(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::I32(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::I64(a), AnyTensor::Bool(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::I64(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::I64(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::U8(a), AnyTensor::Bool(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U8(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::U8(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::U16(a), AnyTensor::Bool(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U16(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::U16(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::U32(a), AnyTensor::Bool(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U32(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::U32(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::U64(a), AnyTensor::Bool(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::U64(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::U64(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::F32(a), AnyTensor::Bool(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F32(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::F32(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::F64(a), AnyTensor::Bool(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::I8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::I16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::I32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::I64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::U8(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::U16(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::U32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::U64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::F32(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::F64(b)) => lift(($f::)($cond, a, b), crate::any_tensor::any_of), + (AnyTensor::F64(a), AnyTensor::C32(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::F64(a), AnyTensor::C64(b)) => { + lift(($f::>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::Bool(b)) => { + lift(($f::, bool>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::I8(b)) => { + lift(($f::, i8>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::I16(b)) => { + lift(($f::, i16>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::I32(b)) => { + lift(($f::, i32>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::I64(b)) => { + lift(($f::, i64>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::U8(b)) => { + lift(($f::, u8>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::U16(b)) => { + lift(($f::, u16>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::U32(b)) => { + lift(($f::, u32>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::U64(b)) => { + lift(($f::, u64>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::F32(b)) => { + lift(($f::, f32>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::F64(b)) => { + lift(($f::, f64>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::C32(b)) => { + lift(($f::, Complex>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C32(a), AnyTensor::C64(b)) => { + lift(($f::, Complex>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::Bool(b)) => { + lift(($f::, bool>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::I8(b)) => { + lift(($f::, i8>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::I16(b)) => { + lift(($f::, i16>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::I32(b)) => { + lift(($f::, i32>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::I64(b)) => { + lift(($f::, i64>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::U8(b)) => { + lift(($f::, u8>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::U16(b)) => { + lift(($f::, u16>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::U32(b)) => { + lift(($f::, u32>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::U64(b)) => { + lift(($f::, u64>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::F32(b)) => { + lift(($f::, f32>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::F64(b)) => { + lift(($f::, f64>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::C32(b)) => { + lift(($f::, Complex>)($cond, a, b), crate::any_tensor::any_of) + }, + (AnyTensor::C64(a), AnyTensor::C64(b)) => { + lift(($f::, Complex>)($cond, a, b), crate::any_tensor::any_of) + }, + } + }; +} +pub(crate) use dispatch_where; + // --------------------------------------------------- typed generic helpers -- pub(crate) fn proj_shape(x: &FTensor) -> Vec diff --git a/crates-interop/rstsr-faer-py/src/lib.rs b/crates-interop/rstsr-faer-py/src/lib.rs index c0b94c26..c8b9ec6f 100644 --- a/crates-interop/rstsr-faer-py/src/lib.rs +++ b/crates-interop/rstsr-faer-py/src/lib.rs @@ -236,6 +236,7 @@ fn rstsr_faer(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(ops::argmin, m)?)?; m.add_function(wrap_pyfunction!(ops::count_nonzero, m)?)?; m.add_function(wrap_pyfunction!(ops::sum_bool, m)?)?; + m.add_function(wrap_pyfunction!(ops::where_, m)?)?; // native name "where" m.add_function(wrap_pyfunction!(indexing::take, m)?)?; // W4 manipulation surface (bindings over rt:: manipulation entries) diff --git a/crates-interop/rstsr-faer-py/src/ops.rs b/crates-interop/rstsr-faer-py/src/ops.rs index 437ba3af..cf57279e 100644 --- a/crates-interop/rstsr-faer-py/src/ops.rs +++ b/crates-interop/rstsr-faer-py/src/ops.rs @@ -35,8 +35,8 @@ use crate::any_tensor::{ device_faer, dispatch_bin_bool_self, dispatch_bin_int_bool_self, dispatch_bin_int_self, dispatch_bin_numeric_self, dispatch_bin_promote, dispatch_bin_promote_eq, dispatch_t, dispatch_t_bool, dispatch_t_float_complex_same, dispatch_t_index_ord, dispatch_t_index_zero, dispatch_t_into_float, dispatch_t_no_complex, dispatch_t_numeric_same, - dispatch_t_real_float_same, dispatch_t_real_numeric_same, dispatch_t_signed, err_py, lift, type_err, AnyTensor, - FTensor, NativeArray, + dispatch_t_real_float_same, dispatch_t_real_numeric_same, dispatch_t_signed, dispatch_where, err_py, lift, + type_err, AnyTensor, FTensor, NativeArray, }; use crate::creation::dim_from; @@ -989,3 +989,35 @@ pub fn sum_bool(x: &NativeArray, axes: Option>, keepdims: bool) -> Py pub fn broadcast_to(x: &NativeArray, shape: Vec) -> PyResult { Ok(NativeArray { t: dispatch_t!(x.t, op_broadcast_to(shape.clone()))? }) } + +// --------------------------------------------------- element-wise select -- +// +// The `where` slice of the searching surface (register G-037). +// +// `where(cond, x, y)`: the condition must be a boolean tensor (spec: "should +// have a boolean data type" — declined, never truthiness-cast); x/y may be any +// dtype pair in rstsr's promotion matrix. The promoted result type comes from +// rstsr's own `DTypePromoteAPI`, so no dtype table is duplicated in the shim. + +fn op_where( + cond: &FTensor, + x: &FTensor, + y: &FTensor, +) -> rt::Result>::Res>> +where + TX: DTypePromoteAPI, + for<'x> &'x FTensor: + TensorWhereAPI<&'x FTensor, &'x FTensor, Output = FTensor<>::Res>>, +{ + rt::where_f(cond, x, y) +} + +/// Rust keyword `where` forces the `where_` item name; the Python-visible name +/// is restored by the pyfunction attribute. +#[pyfunction(name = "where")] +pub fn where_(cond: &NativeArray, x: &NativeArray, y: &NativeArray) -> PyResult { + match &cond.t { + AnyTensor::Bool(c) => Ok(NativeArray { t: dispatch_where!(c, x.t, y.t, op_where)? }), + _ => type_err("where: condition must have a boolean data type"), + } +}