Skip to main content

for_each_batch

Function for_each_batch 

Source
pub unsafe fn for_each_batch<T: Scalar, const N: usize>(
    inputs: [*const TensorView<'_>; N],
    y: *mut TensorViewMut<'_>,
    f: impl FnMut([&[T]; N], &mut [T], usize),
)
Expand description

Call f one time for each lane of the output y of a callback, with the same lane of each input. This is a simple way to write a callback in Rust: f gets one lane of each view as a slice, so it does not use the strides. It is the equivalent of diffsol::Vector::for_each_batch for the views of a callback.

The last argument of f is the index of the lane of y, b * P + k for the lane (k, b) (the lane index of diffsol). The batch lane is lane / P, where P is y.shape()[1].

ⓘ
unsafe extern "C" fn rhs(
    x: *const TensorView<'_>,
    p: *const TensorView<'_>,
    _t: f64,
    y: *mut TensorViewMut<'_>,
    _user_data: *mut c_void,
) -> i32 {
    unsafe { for_each_batch([x, p], y, |[x, p], y: &mut [f64], _lane| y[0] = -p[0] * x[0]) };
    0
}

§Panics

If a view does not have 3 axes, the scalar type of a view is not T, or an input does not have the lanes [P, B] of y.

§Safety

The pointers must be the arguments that the library gave to the callback, and the call must be during the callback.