Skip to main content

erased_select_into_uninit

Function erased_select_into_uninit 

Source
pub fn erased_select_into_uninit(
    dtype: KernelDType,
    ctx: &ExecContext,
    dest: &mut ErasedRawStridedUninitMut<'_>,
    pred: &ErasedRawStridedPtr<'_>,
    on_true: &ErasedRawStridedPtr<'_>,
    on_false: &ErasedRawStridedPtr<'_>,
) -> Result<()>
Expand description

Select elementwise between two same-dtype operands by a bool predicate, writing uninitialized storage: dest[i] = if pred[i] { on_true[i] } else { on_false[i] }.

Every dtype is supported. The predicate dtype must be bool.

§Examples

use core::mem::MaybeUninit;
use strided_kernel::{
    erased_select_into_uninit, ErasedRawStridedPtr, ErasedRawStridedRef,
    ErasedRawStridedUninitMut, ExecContext, KernelDType,
};

let pred = [true, false];
let on_true = [1_i64, 2];
let on_false = [10_i64, 20];
let pred = ErasedRawStridedRef::from_slice(&pred, &[2], &[1], 0).unwrap();
let on_true = ErasedRawStridedRef::from_slice(&on_true, &[2], &[1], 0).unwrap();
let on_false = ErasedRawStridedRef::from_slice(&on_false, &[2], &[1], 0).unwrap();
let mut out = [MaybeUninit::<i64>::uninit(); 2];
let mut dest = ErasedRawStridedUninitMut::from_uninit_slice(&mut out, &[2], &[1], 0).unwrap();
erased_select_into_uninit(
    KernelDType::I64,
    &ExecContext::serial(),
    &mut dest,
    &ErasedRawStridedPtr::from_ref(&pred),
    &ErasedRawStridedPtr::from_ref(&on_true),
    &ErasedRawStridedPtr::from_ref(&on_false),
)
.unwrap();
// SAFETY: a successful call initializes every reachable element.
assert_eq!(unsafe { [out[0].assume_init(), out[1].assume_init()] }, [1, 20]);

§Errors

Returns a typed StridedError for dtype, shape, output-layout, or overlap contracts. Validation completes before any write.