Skip to main content

for_lanes

Function for_lanes 

Source
pub fn for_lanes<Op: LaneOp>(lanes: usize, op: Op) -> Op::Out
Expand description

Call f monomorphised at the width width_for picks for lanes.

The const parameter has to be known at compile time, so the runtime column count is turned into one of a fixed set of instantiations here rather than at every call site.

use phyz_diff::multidual::{MultiDual, for_lanes};
use tang::Scalar;

// Differentiate f(x) = Σ xᵢ² in however many directions were asked for.
let xs = [1.0, 2.0, 3.0];
let grad = for_lanes(xs.len(), Grad(&xs));
assert_eq!(grad, vec![2.0, 4.0, 6.0]);

struct Grad<'a>(&'a [f64]);
impl phyz_diff::multidual::LaneOp for Grad<'_> {
    type Out = Vec<f64>;
    fn call<const N: usize>(self) -> Vec<f64> {
        let seeded: Vec<MultiDual<N>> = self
            .0
            .iter()
            .enumerate()
            .map(|(i, &x)| MultiDual::var(x, i))
            .collect();
        let f = seeded
            .iter()
            .fold(MultiDual::<N>::constant(0.0), |acc, &x| acc + x * x);
        (0..self.0.len()).map(|i| f.dual[i]).collect()
    }
}