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"), + } +} 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 */