nuvec_core/
fused.rs

1//! Fused element-wise computations: one kernel for a graph of element-wise operations
2//! (MLX's `Compiled` primitive), as built by `compile`.
3//!
4//! On the GPU the graph becomes one generated Metal kernel that reads each input once and
5//! writes only the result. Everywhere else (and for the transforms) it expands into the
6//! operations it fuses.
7
8use std::fmt::Write as _;
9use std::hash::{DefaultHasher, Hash, Hasher};
10use std::sync::Arc;
11
12use crate::array::Data;
13use crate::kernels::Operand;
14use crate::primitive::{BinaryOp, Primitive, UnaryOp};
15use crate::{Array, Category, DType, Error, Result};
16
17/// An operand of a fused step: an input of the fused computation or an earlier step's
18/// result.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20pub enum FusedRef {
21    /// The input at this position.
22    Input(usize),
23    /// The result of the step at this position.
24    Step(usize),
25    /// The constant at this position, written into the kernel as a literal (as MLX's
26    /// compiled kernels write theirs).
27    Constant(usize),
28}
29
30/// A graph of element-wise operations evaluated as one kernel, with one or more of its
31/// steps as results. Cheap to clone; equal graphs compare and hash equal.
32#[derive(Debug, Clone, PartialEq, Eq, Hash)]
33pub struct Fused(Arc<Graph>);
34
35#[derive(Debug, PartialEq, Eq, Hash)]
36struct Graph {
37    inputs: Vec<DType>,
38    /// Scalar constants: a float dtype and the value's bits.
39    constants: Vec<(DType, u64)>,
40    steps: Vec<(Primitive, Vec<FusedRef>)>,
41    /// The steps whose results are outputs.
42    outputs: Vec<usize>,
43    /// The element type of each step's result.
44    dtypes: Vec<DType>,
45    /// A hash of the above, naming the generated kernels.
46    id: u64,
47}
48
49/// How a generated kernel reads one input.
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51pub(crate) enum Load {
52    /// One element, broadcast.
53    Scalar,
54    /// Laid out like the output.
55    Flat,
56    /// Through shape and strides (the `k`-th strided input).
57    Strided(usize),
58}
59
60fn computes(d: DType) -> bool {
61    !matches!(d, DType::F64 | DType::Complex64)
62}
63
64fn is_float(d: DType) -> bool {
65    matches!(d, DType::F16 | DType::BF16 | DType::F32)
66}
67
68impl Fused {
69    /// Whether `primitive` applied to inputs of these dtypes can be part of a fused
70    /// kernel: unary, binary, select and type conversions over the GPU's element types
71    /// (all but float64 and complex64).
72    pub fn supports(primitive: &Primitive, dtypes: &[DType]) -> bool {
73        if !dtypes.iter().copied().all(computes) {
74            return false;
75        }
76        match (primitive, dtypes) {
77            (Primitive::Unary(op), &[d]) => match op {
78                UnaryOp::Negative => d != DType::Bool,
79                UnaryOp::Abs => true,
80                UnaryOp::Conjugate => false,
81                _ => is_float(d),
82            },
83            (Primitive::Binary(op), &[a, b]) => {
84                a == b && (!matches!(op, BinaryOp::Divide | BinaryOp::LogAddExp) || is_float(a))
85            }
86            (Primitive::Select, &[c, a, b]) => c == DType::Bool && a == b,
87            (Primitive::AsType(to), &[_]) => computes(*to),
88            _ => false,
89        }
90    }
91
92    /// The fused computation of `steps` over inputs of the given dtypes, with the results
93    /// of the steps at positions `outputs` as its outputs. Each step is a primitive
94    /// [`Fused::supports`] and its operands, which refer to inputs or earlier steps.
95    pub fn new(
96        inputs: &[DType],
97        steps: Vec<(Primitive, Vec<FusedRef>)>,
98        outputs: &[usize],
99    ) -> Result<Self> {
100        Self::with_constants(inputs, &[], steps, outputs)
101    }
102
103    /// [`Fused::new`] with scalar `constants` (a float dtype, `F32`, `F16` or `BF16`, and the
104    /// value's bits) that steps refer to as [`FusedRef::Constant`]: they are part of the
105    /// kernel, as literals, rather than inputs read from memory.
106    pub fn with_constants(
107        inputs: &[DType],
108        constants: &[(DType, u64)],
109        steps: Vec<(Primitive, Vec<FusedRef>)>,
110        outputs: &[usize],
111    ) -> Result<Self> {
112        let invalid = |what: String| Err(Error::InvalidArgument(format!("fused: {what}")));
113        if steps.is_empty() || outputs.is_empty() {
114            return invalid("no steps or no outputs".into());
115        }
116        if let Some(o) = outputs.iter().find(|&&o| o >= steps.len()) {
117            return invalid(format!("output step {o} of {}", steps.len()));
118        }
119        if let Some((d, _)) = constants
120            .iter()
121            .find(|(d, _)| !matches!(d, DType::F32 | DType::F16 | DType::BF16))
122        {
123            return invalid(format!("a {d:?} constant"));
124        }
125        let mut dtypes: Vec<DType> = Vec::with_capacity(steps.len());
126        for (k, (primitive, args)) in steps.iter().enumerate() {
127            let mut ds = Vec::with_capacity(args.len());
128            for arg in args {
129                ds.push(match *arg {
130                    FusedRef::Input(i) if i < inputs.len() => inputs[i],
131                    FusedRef::Step(j) if j < k => dtypes[j],
132                    FusedRef::Constant(c) if c < constants.len() => constants[c].0,
133                    _ => return invalid(format!("step {k} refers to {arg:?}")),
134                });
135            }
136            if !Self::supports(primitive, &ds) {
137                return invalid(format!("{} on {ds:?} cannot be fused", primitive.name()));
138            }
139            dtypes.push(match primitive {
140                Primitive::Binary(op) if op.is_comparison() => DType::Bool,
141                Primitive::Select => ds[1],
142                Primitive::AsType(to) => *to,
143                _ => ds[0],
144            });
145        }
146        let mut graph = Graph {
147            inputs: inputs.to_vec(),
148            constants: constants.to_vec(),
149            steps,
150            outputs: outputs.to_vec(),
151            dtypes,
152            id: 0,
153        };
154        let mut h = DefaultHasher::new();
155        (
156            &graph.inputs,
157            &graph.constants,
158            &graph.steps,
159            &graph.outputs,
160        )
161            .hash(&mut h);
162        graph.id = h.finish();
163        Ok(Self(Arc::new(graph)))
164    }
165
166    /// The input dtypes.
167    pub fn inputs(&self) -> &[DType] {
168        &self.0.inputs
169    }
170
171    /// The number of fused operations.
172    pub fn len(&self) -> usize {
173        self.0.steps.len()
174    }
175
176    /// Always false: a fused computation has at least one step.
177    pub fn is_empty(&self) -> bool {
178        false
179    }
180
181    /// The outputs' dtypes.
182    pub fn dtypes(&self) -> Vec<DType> {
183        self.0.outputs.iter().map(|&o| self.0.dtypes[o]).collect()
184    }
185
186    /// The same computation recorded as ordinary operations.
187    pub(crate) fn expand(&self, inputs: &[Array]) -> Result<Vec<Array>> {
188        let mut results: Vec<Array> = Vec::with_capacity(self.0.steps.len());
189        let shape = inputs.first().map_or(&[][..], Array::dims);
190        for (primitive, args) in &self.0.steps {
191            let args = args
192                .iter()
193                .map(|r| match *r {
194                    FusedRef::Input(i) => Ok(inputs[i].clone()),
195                    FusedRef::Step(j) => Ok(results[j].clone()),
196                    // The steps take operands of one shape, the inputs' (the kernel writes
197                    // the constant as a literal).
198                    FusedRef::Constant(c) => {
199                        crate::ops::imp::broadcast_to(&constant_array(self.0.constants[c]), shape)
200                    }
201                })
202                .collect::<Result<_>>()?;
203            results.push(Array::from_op(primitive.clone(), args)?.swap_remove(0));
204        }
205        Ok(self.0.outputs.iter().map(|&o| results[o].clone()).collect())
206    }
207
208    /// Evaluates the expanded operations on the CPU, for work the GPU does not take.
209    pub(crate) fn eval_cpu(&self, operands: &[Operand<'_>]) -> Result<Vec<Data>> {
210        crate::fast::imp::run_on_cpu(operands, |leaves| self.expand(leaves))
211    }
212
213    /// Reverse mode through the expanded operations, for the wanted float inputs.
214    pub(crate) fn vjp(
215        &self,
216        inputs: &[Array],
217        cotangents: &[Array],
218        wanted: &[bool],
219    ) -> Result<Vec<Option<Array>>> {
220        let active: Vec<usize> = (0..inputs.len())
221            .filter(|&i| wanted[i] && inputs[i].dtype().category() == Category::Float)
222            .collect();
223        let primals: Vec<Array> = active.iter().map(|&i| inputs[i].clone()).collect();
224        let f = |ps: &[Array]| self.expand(&substitute(inputs, &active, ps));
225        let (_, grads) = crate::vjp_flat(f, &primals, cotangents)?;
226        let mut out = vec![None; inputs.len()];
227        for (g, &i) in grads.into_iter().zip(&active) {
228            out[i] = Some(g);
229        }
230        Ok(out)
231    }
232
233    /// Forward mode through the expanded operations (non-float inputs are constants).
234    pub(crate) fn jvp(&self, inputs: &[Array], tangents: &[Option<Array>]) -> Result<Vec<Array>> {
235        let active: Vec<usize> = (0..inputs.len())
236            .filter(|&i| tangents[i].is_some() && inputs[i].dtype().category() == Category::Float)
237            .collect();
238        let primals: Vec<Array> = active.iter().map(|&i| inputs[i].clone()).collect();
239        let ts: Vec<Array> = active
240            .iter()
241            .map(|&i| tangents[i].clone().expect("active"))
242            .collect();
243        if active.is_empty() {
244            return self
245                .expand(inputs)?
246                .iter()
247                .map(crate::derivatives::zeros_like)
248                .collect();
249        }
250        let f = |ps: &[Array]| self.expand(&substitute(inputs, &active, ps));
251        Ok(crate::jvp_flat(f, &primals, &ts)?.1)
252    }
253
254    /// Batching rule: the expanded operations, batched.
255    pub(crate) fn vmap(
256        &self,
257        inputs: &[Array],
258        axes: &[Option<usize>],
259    ) -> Result<(Vec<Array>, usize)> {
260        let f = |xs: &[Array]| self.expand(xs);
261        Ok((crate::transforms::batch(&f, inputs, axes)?, 0))
262    }
263
264    /// The kernel name for inputs read as `loads`, with `per_thread` elements per thread.
265    pub(crate) fn kernel_name(&self, loads: &[Load], per_thread: usize) -> String {
266        let mut name = format!("nuvec_fused_{:016x}_", self.0.id);
267        name.extend(loads.iter().map(|l| match l {
268            Load::Scalar => 's',
269            Load::Flat => 'f',
270            Load::Strided(_) => 'g',
271        }));
272        let _ = write!(name, "_n{per_thread}");
273        name
274    }
275
276    /// Metal source of the kernel `name`, reading inputs as `loads`, `per_thread`
277    /// consecutive elements per thread. Arguments: the inputs, the outputs, the inputs'
278    /// element offsets, the (collapsed) shape, the strided inputs' strides (`ndim` each,
279    /// in order), `ndim` and the element count.
280    pub(crate) fn metal_source(&self, name: &str, loads: &[Load], per_thread: usize) -> String {
281        let g = &*self.0;
282        let (n, m) = (g.inputs.len(), g.outputs.len());
283        let mut s = String::from(
284            "#include \"common.h\"\nusing namespace metal;\nusing namespace nuvec;\n\n",
285        );
286        let _ = writeln!(s, "[[kernel]] void {name}(");
287        for (i, d) in g.inputs.iter().enumerate() {
288            let _ = writeln!(s, "    device const {}* in{i} [[buffer({i})]],", msl(*d));
289        }
290        for (k, d) in self.dtypes().into_iter().enumerate() {
291            let _ = writeln!(s, "    device {}* out{k} [[buffer({})]],", msl(d), n + k);
292        }
293        let _ = writeln!(s, "    constant long* offs [[buffer({})]],", n + m);
294        let _ = writeln!(s, "    constant int* shape [[buffer({})]],", n + m + 1);
295        let _ = writeln!(s, "    constant long* strides [[buffer({})]],", n + m + 2);
296        let _ = writeln!(s, "    constant int& ndim [[buffer({})]],", n + m + 3);
297        let _ = writeln!(s, "    constant uint& size [[buffer({})]],", n + m + 4);
298        s.push_str("    uint tid [[thread_position_in_grid]]) {\n");
299        for (i, d) in g.inputs.iter().enumerate() {
300            let _ = writeln!(s, "  device const {}* a{i} = in{i} + offs[{i}];", msl(*d));
301        }
302        // Constants as float literals, as MLX writes them: the compiler folds them early (a
303        // literal exponent specializes `pow`). The shortest round-trip form reads back as the
304        // same float, which holds half and bfloat values exactly; non-finite values are
305        // written from their bits.
306        for (c, &(d, bits)) in g.constants.iter().enumerate() {
307            let value = match d {
308                DType::F16 => crate::f16::from_bits(bits as u16).to_f32(),
309                DType::BF16 => crate::bf16::from_bits(bits as u16).to_f32(),
310                _ => f32::from_bits(bits as u32),
311            };
312            let t = msl(d);
313            if value.is_finite() {
314                let _ = writeln!(s, "  const {t} c{c} = {t}({value:?}f);");
315            } else {
316                let _ = writeln!(
317                    s,
318                    "  const {t} c{c} = {t}(as_type<float>(0x{:x}u));",
319                    value.to_bits()
320                );
321            }
322        }
323        let _ = writeln!(
324            s,
325            "  for (uint i = tid * {per_thread}; i < min(size, (tid + 1) * {per_thread}); ++i) {{"
326        );
327        for (i, (d, load)) in g.inputs.iter().zip(loads).enumerate() {
328            let index = match load {
329                Load::Scalar => "0".into(),
330                Load::Flat => "i".into(),
331                Load::Strided(k) => format!("loc(i, shape, strides + {k} * ndim, ndim)"),
332            };
333            let _ = writeln!(s, "    {} x{i} = a{i}[{index}];", msl(*d));
334        }
335        let arg = |r: &FusedRef| match *r {
336            FusedRef::Input(i) => format!("x{i}"),
337            FusedRef::Step(j) => format!("t{j}"),
338            FusedRef::Constant(c) => format!("c{c}"),
339        };
340        let dtype = |r: &FusedRef| match *r {
341            FusedRef::Input(i) => g.inputs[i],
342            FusedRef::Step(j) => g.dtypes[j],
343            FusedRef::Constant(c) => g.constants[c].0,
344        };
345        for (k, (primitive, args)) in g.steps.iter().enumerate() {
346            let a: Vec<String> = args.iter().map(arg).collect();
347            let expr = match primitive {
348                Primitive::Unary(op) => format!("op_{}({})", unary_fn(*op), a[0]),
349                Primitive::Binary(op) => binary_expr(*op, &a[0], &a[1]),
350                Primitive::Select => format!("({}.v ? {} : {})", a[0], a[1], a[2]),
351                Primitive::AsType(to) => {
352                    format!("convert<{}, {}>({})", msl(dtype(&args[0])), msl(*to), a[0])
353                }
354                _ => unreachable!("validated by Fused::new"),
355            };
356            let _ = writeln!(s, "    {} t{k} = {expr};", msl(g.dtypes[k]));
357        }
358        for (k, o) in g.outputs.iter().enumerate() {
359            let _ = writeln!(s, "    out{k}[i] = t{o};");
360        }
361        s.push_str("  }\n}\n");
362        s
363    }
364}
365
366/// A fused constant as a scalar array (for the composition).
367fn constant_array((dtype, bits): (DType, u64)) -> Array {
368    match dtype {
369        DType::F16 => Array::from(crate::f16::from_bits(bits as u16)),
370        DType::BF16 => Array::from(crate::bf16::from_bits(bits as u16)),
371        _ => Array::from(f32::from_bits(bits as u32)),
372    }
373}
374
375/// `inputs` with the positions in `active` replaced by `values`.
376fn substitute(inputs: &[Array], active: &[usize], values: &[Array]) -> Vec<Array> {
377    let mut all = inputs.to_vec();
378    for (v, &i) in values.iter().zip(active) {
379        all[i] = v.clone();
380    }
381    all
382}
383
384/// The Metal type of a GPU dtype (see `common.h`).
385fn msl(d: DType) -> &'static str {
386    match d {
387        DType::Bool => "nbool",
388        DType::U8 => "uchar",
389        DType::U16 => "ushort",
390        DType::U32 => "uint",
391        DType::U64 => "ulong",
392        DType::I8 => "char",
393        DType::I16 => "short",
394        DType::I32 => "int",
395        DType::I64 => "long",
396        DType::F16 => "half",
397        DType::BF16 => "bfloat",
398        DType::F32 => "float",
399        DType::F64 | DType::Complex64 => unreachable!("not fused"),
400    }
401}
402
403fn unary_fn(op: UnaryOp) -> &'static str {
404    match op {
405        UnaryOp::Negative => "neg",
406        UnaryOp::Abs => "abs",
407        UnaryOp::Exp => "exp",
408        UnaryOp::Log => "log",
409        UnaryOp::Sqrt => "sqrt",
410        UnaryOp::Sin => "sin",
411        UnaryOp::Cos => "cos",
412        UnaryOp::Tanh => "tanh",
413        UnaryOp::Sigmoid => "sigmoid",
414        UnaryOp::ErfInv => "erfinv",
415        UnaryOp::Conjugate => unreachable!("not fused"),
416    }
417}
418
419/// The element-wise kernels' expressions (`elementwise.metal`).
420fn binary_expr(op: BinaryOp, a: &str, b: &str) -> String {
421    match op {
422        BinaryOp::Add => format!("op_add({a}, {b})"),
423        BinaryOp::Subtract => format!("op_sub({a}, {b})"),
424        BinaryOp::Multiply => format!("op_mul({a}, {b})"),
425        BinaryOp::Divide => format!("op_div({a}, {b})"),
426        BinaryOp::Maximum => format!("op_max({a}, {b})"),
427        BinaryOp::Minimum => format!("op_min({a}, {b})"),
428        BinaryOp::LogAddExp => format!("op_logaddexp({a}, {b})"),
429        BinaryOp::FloorDivide => format!("op_floordiv({a}, {b})"),
430        BinaryOp::Power => format!("op_pow({a}, {b})"),
431        BinaryOp::Equal => format!("nbool{{uchar(op_eq({a}, {b}))}}"),
432        BinaryOp::NotEqual => format!("nbool{{uchar(!op_eq({a}, {b}))}}"),
433        BinaryOp::Less => format!("nbool{{uchar(op_lt({a}, {b}))}}"),
434        BinaryOp::LessEqual => format!("nbool{{uchar(op_lt({a}, {b}) || op_eq({a}, {b}))}}"),
435        BinaryOp::Greater => format!("nbool{{uchar(op_lt({b}, {a}))}}"),
436        BinaryOp::GreaterEqual => format!("nbool{{uchar(op_lt({b}, {a}) || op_eq({a}, {b}))}}"),
437    }
438}
439
440#[cfg(test)]
441mod tests {
442    use super::*;
443    use crate::ops::imp as ops;
444
445    fn sigmoid_mul() -> Fused {
446        // x * sigmoid(x) + y
447        Fused::new(
448            &[DType::F32, DType::F32],
449            vec![
450                (Primitive::Unary(UnaryOp::Sigmoid), vec![FusedRef::Input(0)]),
451                (
452                    Primitive::Binary(BinaryOp::Multiply),
453                    vec![FusedRef::Input(0), FusedRef::Step(0)],
454                ),
455                (
456                    Primitive::Binary(BinaryOp::Add),
457                    vec![FusedRef::Step(1), FusedRef::Input(1)],
458                ),
459            ],
460            &[1, 2],
461        )
462        .unwrap()
463    }
464
465    #[test]
466    fn validates_steps() {
467        let bad = |inputs: &[DType], steps| Fused::new(inputs, steps, &[0]).is_err();
468        assert!(bad(&[DType::F32], vec![]));
469        assert!(bad(
470            &[DType::F32],
471            vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Step(0)])]
472        ));
473        assert!(bad(
474            &[DType::I32],
475            vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Input(0)])]
476        ));
477        assert!(bad(
478            &[DType::F64],
479            vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Input(0)])]
480        ));
481        assert!(bad(&[DType::F32], vec![(Primitive::Matmul, vec![])]));
482        let f = sigmoid_mul();
483        assert_eq!((f.len(), f.dtypes()), (3, vec![DType::F32; 2]));
484        let one = vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Input(0)])];
485        assert!(
486            Fused::new(&[DType::F32], one, &[1]).is_err(),
487            "no such output step"
488        );
489        assert_eq!(f, sigmoid_mul(), "equal graphs are equal");
490    }
491
492    fn values(arrays: &[Array]) -> Vec<Vec<f32>> {
493        arrays.iter().map(|a| a.to_vec::<f32>().unwrap()).collect()
494    }
495
496    /// Equal up to rounding (Miri perturbs transcendental functions by a few ulps).
497    fn assert_close(a: &[Vec<f32>], b: &[Vec<f32>]) {
498        assert_eq!(a.len(), b.len());
499        for (x, y) in a.iter().flatten().zip(b.iter().flatten()) {
500            assert!((x - y).abs() <= 1e-5 * x.abs().max(1.0), "{a:?} vs {b:?}");
501        }
502    }
503
504    #[test]
505    fn evaluates_like_its_operations_on_the_cpu() {
506        crate::set_default_stream(crate::Stream::default_for(crate::Device::cpu()));
507        let x = Array::from_slice(&[-2.0f32, 0.5, 3.0, 1.0, -1.0, 4.0], [2, 3]);
508        let y = Array::from_slice(&[1.0f32, 2.0], [2, 1]);
509        let y = ops::broadcast_to(&y, &[2, 3]).unwrap();
510        let xt = x.t().t();
511        let fused = sigmoid_mul();
512        let out = Array::from_op(Primitive::Compiled(fused.clone()), vec![xt, y.clone()]).unwrap();
513        assert_eq!(out.len(), 2);
514        assert_close(&values(&out), &values(&fused.expand(&[x, y]).unwrap()));
515    }
516
517    #[test]
518    fn expands_its_constants_to_the_inputs_shape() {
519        crate::set_default_stream(crate::Stream::default_for(crate::Device::cpu()));
520        // x * 0.5 + x
521        let fused = Fused::with_constants(
522            &[DType::BF16],
523            &[(DType::BF16, u64::from(crate::bf16::from_f32(0.5).to_bits()))],
524            vec![
525                (
526                    Primitive::Binary(BinaryOp::Multiply),
527                    vec![FusedRef::Input(0), FusedRef::Constant(0)],
528                ),
529                (
530                    Primitive::Binary(BinaryOp::Add),
531                    vec![FusedRef::Step(0), FusedRef::Input(0)],
532                ),
533            ],
534            &[1],
535        )
536        .unwrap();
537        let x = Array::from_slice(&[2.0f32, -4.0, 1.0, 6.0], [2, 2]).astype(DType::BF16);
538        let f = |xs: &[Array]| Array::from_op(Primitive::Compiled(fused.clone()), xs.to_vec());
539        let out = f(std::slice::from_ref(&x)).unwrap()[0].astype(DType::F32);
540        assert_eq!(out.to_vec::<f32>().unwrap(), [3.0, -6.0, 1.5, 9.0]);
541        let ones = crate::ops::ones_like(&x);
542        let (_, g) = crate::vjp_flat(f, &[x], &[ones]).unwrap();
543        assert_eq!(g[0].astype(DType::F32).to_vec::<f32>().unwrap(), [1.5; 4]);
544    }
545
546    #[test]
547    fn differentiates_through_its_operations() {
548        crate::set_default_stream(crate::Stream::default_for(crate::Device::cpu()));
549        let fused = sigmoid_mul();
550        let x = Array::from_slice(&[-1.0f32, 0.0, 2.0], [3]);
551        let y = Array::from_slice(&[3.0f32, 1.0, 0.0], [3]);
552        let cot = Array::from_slice(&[1.0f32, 2.0, 0.5], [3]);
553        let cot2 = Array::from_slice(&[-1.0f32, 0.25, 3.0], [3]);
554        let run = |expanded: bool| {
555            let f = |xs: &[Array]| match expanded {
556                true => fused.expand(xs),
557                false => Array::from_op(Primitive::Compiled(fused.clone()), xs.to_vec()),
558            };
559            let primals = [x.clone(), y.clone()];
560            let (_, g) = crate::vjp_flat(f, &primals, &[cot.clone(), cot2.clone()]).unwrap();
561            let (_, t) = crate::jvp_flat(f, &primals, &[cot.clone(), cot2.clone()]).unwrap();
562            values(&[g, t].concat())
563        };
564        assert_close(&run(false), &run(true));
565        let batched = crate::vmap_flat(
566            |xs| Array::from_op(Primitive::Compiled(fused.clone()), xs.to_vec()),
567            &[Some(0), None],
568            &[],
569        );
570        let xs = ops::reshape(&x, &[3, 1]).unwrap();
571        let ys = Array::from_slice(&[3.0f32], [1]);
572        let out = batched(&[xs.clone(), ys.clone()]).unwrap();
573        let expected = fused
574            .expand(&[xs, ops::broadcast_to(&ys, &[3, 1]).unwrap()])
575            .unwrap();
576        assert_close(&values(&out), &values(&expected));
577    }
578}
nuvec_core/fused.rs

nuvec_core/
fused.rs
