rlx_cpu/
spd_kernels.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! CPU kernels for the core Riemannian / SPD-manifold ops (`Op::BiMap`,
6//! `Op::ReEig`, `Op::LogEig`, `Op::SpdBatchNorm`, `Op::SpdKarcherMean`
7//! and their backwards).
8//!
9//! These reuse the `CpuKernel` scaffold — the same typed-view contract
10//! as user custom ops — so the core ops ride the existing
11//! `Thunk::CustomOp` execution path (both the closure-compiled and the
12//! match-based executors) without bespoke `Thunk` variants. The compute
13//! bodies live in [`crate::spd`]; per-op parameters (`eps`/`iters`/`tol`)
14//! are baked into the kernel struct at compile time (F64, CPU-first).
15
16use crate::op_registry::{CpuKernel, CpuTensorMut, CpuTensorRef};
17use crate::spd;
18use rlx_ir::SpdMatFn;
19
20/// Read the `(rows, cols)` of an `[n,n]`-style square matrix input.
21fn dim(t: &CpuTensorRef<'_>, i: usize) -> usize {
22    t.shape().dim(i).unwrap_static()
23}
24
25/// BiMap forward: `Y = W · X · Wᵀ`. Inputs `[W [m,n], X [n,n]]`.
26pub struct BiMapKernel;
27impl CpuKernel for BiMapKernel {
28    fn name(&self) -> &str {
29        "core.bimap"
30    }
31    fn execute(
32        &self,
33        inputs: &[CpuTensorRef<'_>],
34        output: CpuTensorMut<'_>,
35        _attrs: &[u8],
36    ) -> Result<(), String> {
37        let w = inputs[0].expect_f64("bimap W")?;
38        let x = inputs[1].expect_f64("bimap X")?;
39        let (m, n) = (dim(&inputs[0], 0), dim(&inputs[0], 1));
40        let out = output.expect_f64_mut("bimap Y")?;
41        out.copy_from_slice(&spd::bimap(w, x, m, n));
42        Ok(())
43    }
44}
45
46/// ReEig forward: `Y = U max(ε, Σ) Uᵀ`. Input `[X [n,n]]`; output is
47/// the packed `[Y (n²), λ (n), U (n²)]` so the backward reuses the
48/// eigendecomposition. The builder narrows out `Y`.
49pub struct ReEigKernel {
50    pub eps: f64,
51}
52impl CpuKernel for ReEigKernel {
53    fn name(&self) -> &str {
54        "core.reeig"
55    }
56    fn execute(
57        &self,
58        inputs: &[CpuTensorRef<'_>],
59        output: CpuTensorMut<'_>,
60        _attrs: &[u8],
61    ) -> Result<(), String> {
62        let x = inputs[0].expect_f64("reeig X")?;
63        let n = dim(&inputs[0], 0);
64        let eps = self.eps;
65        let out = output.expect_f64_mut("reeig packed")?;
66        out.copy_from_slice(&spd::spectral_forward_packed(x, n, |l| l.max(eps)));
67        Ok(())
68    }
69}
70
71/// LogEig forward: `Y = logm(X)`. Packed `[Y, λ, U]` output (see [`ReEigKernel`]).
72pub struct LogEigKernel {
73    pub eps: f64,
74}
75impl CpuKernel for LogEigKernel {
76    fn name(&self) -> &str {
77        "core.logeig"
78    }
79    fn execute(
80        &self,
81        inputs: &[CpuTensorRef<'_>],
82        output: CpuTensorMut<'_>,
83        _attrs: &[u8],
84    ) -> Result<(), String> {
85        let x = inputs[0].expect_f64("logeig X")?;
86        let n = dim(&inputs[0], 0);
87        let eps = self.eps;
88        let out = output.expect_f64_mut("logeig packed")?;
89        out.copy_from_slice(&spd::spectral_forward_packed(x, n, |l| l.max(eps).ln()));
90        Ok(())
91    }
92}
93
94/// SPD batch-norm transport forward. Inputs `[X [B,n,n], mean [n,n], G [n,n]]`.
95pub struct SpdBatchNormKernel {
96    pub eps: f64,
97}
98impl CpuKernel for SpdBatchNormKernel {
99    fn name(&self) -> &str {
100        "core.spd_batch_norm"
101    }
102    fn execute(
103        &self,
104        inputs: &[CpuTensorRef<'_>],
105        output: CpuTensorMut<'_>,
106        _attrs: &[u8],
107    ) -> Result<(), String> {
108        let x = inputs[0].expect_f64("spd_bn X")?;
109        let mean = inputs[1].expect_f64("spd_bn mean")?;
110        let g = inputs[2].expect_f64("spd_bn G")?;
111        let (batch, n) = (dim(&inputs[0], 0), dim(&inputs[0], 1));
112        let out = output.expect_f64_mut("spd_bn Y")?;
113        out.copy_from_slice(&spd::spd_bn_transport(x, mean, g, batch, n, self.eps));
114        Ok(())
115    }
116}
117
118/// Batch Fréchet/Karcher mean. Input `[X [B,n,n]]` → `[n,n]`.
119pub struct SpdKarcherMeanKernel {
120    pub iters: usize,
121    pub tol: f64,
122}
123impl CpuKernel for SpdKarcherMeanKernel {
124    fn name(&self) -> &str {
125        "core.spd_karcher_mean"
126    }
127    fn execute(
128        &self,
129        inputs: &[CpuTensorRef<'_>],
130        output: CpuTensorMut<'_>,
131        _attrs: &[u8],
132    ) -> Result<(), String> {
133        let x = inputs[0].expect_f64("karcher X")?;
134        let (batch, n) = (dim(&inputs[0], 0), dim(&inputs[0], 1));
135        let covs: Vec<Vec<f64>> = (0..batch)
136            .map(|bi| x[bi * n * n..(bi + 1) * n * n].to_vec())
137            .collect();
138        let out = output.expect_f64_mut("karcher mean")?;
139        out.copy_from_slice(&spd::karcher_mean(&covs, n, self.iters, self.tol));
140        Ok(())
141    }
142}
143
144/// Weighted batch Fréchet/Karcher mean. Inputs `[X [B,n,n], w [B]]` → `[n,n]`.
145pub struct SpdKarcherMeanWeightedKernel {
146    pub iters: usize,
147    pub tol: f64,
148}
149impl CpuKernel for SpdKarcherMeanWeightedKernel {
150    fn name(&self) -> &str {
151        "core.spd_karcher_mean_weighted"
152    }
153    fn execute(
154        &self,
155        inputs: &[CpuTensorRef<'_>],
156        output: CpuTensorMut<'_>,
157        _attrs: &[u8],
158    ) -> Result<(), String> {
159        let x = inputs[0].expect_f64("karcher_w X")?;
160        let w = inputs[1].expect_f64("karcher_w weights")?;
161        let (batch, n) = (dim(&inputs[0], 0), dim(&inputs[0], 1));
162        let covs: Vec<Vec<f64>> = (0..batch)
163            .map(|bi| x[bi * n * n..(bi + 1) * n * n].to_vec())
164            .collect();
165        let out = output.expect_f64_mut("karcher_w mean")?;
166        out.copy_from_slice(&spd::karcher_mean_weighted(
167            &covs, w, n, self.iters, self.tol,
168        ));
169        Ok(())
170    }
171}
172
173/// AIRM log map at an arbitrary base. Inputs `[base [n,n], X [n,n]]` → `[n,n]`.
174pub struct SpdLogMapKernel;
175impl CpuKernel for SpdLogMapKernel {
176    fn name(&self) -> &str {
177        "core.spd_log_map"
178    }
179    fn execute(
180        &self,
181        inputs: &[CpuTensorRef<'_>],
182        output: CpuTensorMut<'_>,
183        _attrs: &[u8],
184    ) -> Result<(), String> {
185        let base = inputs[0].expect_f64("log_map base")?;
186        let x = inputs[1].expect_f64("log_map X")?;
187        let n = dim(&inputs[0], 0);
188        let out = output.expect_f64_mut("log_map V")?;
189        out.copy_from_slice(&spd::log_map(base, x, n));
190        Ok(())
191    }
192}
193
194/// AIRM exp map at an arbitrary base. Inputs `[base [n,n], V [n,n]]` → `[n,n]`.
195pub struct SpdExpMapKernel;
196impl CpuKernel for SpdExpMapKernel {
197    fn name(&self) -> &str {
198        "core.spd_exp_map"
199    }
200    fn execute(
201        &self,
202        inputs: &[CpuTensorRef<'_>],
203        output: CpuTensorMut<'_>,
204        _attrs: &[u8],
205    ) -> Result<(), String> {
206        let base = inputs[0].expect_f64("exp_map base")?;
207        let v = inputs[1].expect_f64("exp_map V")?;
208        let n = dim(&inputs[0], 0);
209        let out = output.expect_f64_mut("exp_map Y")?;
210        out.copy_from_slice(&spd::exp_map(base, v, n));
211        Ok(())
212    }
213}
214
215/// AIRM parallel transport. Inputs `[from [n,n], to [n,n], V [n,n]]` → `[n,n]`.
216pub struct SpdParallelTransportKernel;
217impl CpuKernel for SpdParallelTransportKernel {
218    fn name(&self) -> &str {
219        "core.spd_parallel_transport"
220    }
221    fn execute(
222        &self,
223        inputs: &[CpuTensorRef<'_>],
224        output: CpuTensorMut<'_>,
225        _attrs: &[u8],
226    ) -> Result<(), String> {
227        let from = inputs[0].expect_f64("transport from")?;
228        let to = inputs[1].expect_f64("transport to")?;
229        let v = inputs[2].expect_f64("transport V")?;
230        let n = dim(&inputs[0], 0);
231        let out = output.expect_f64_mut("transport out")?;
232        out.copy_from_slice(&spd::parallel_transport(from, to, v, n));
233        Ok(())
234    }
235}
236
237/// Batched SPD spectral matrix function. Input `[X [B,n,n]]` → `[B,n,n]`.
238pub struct SpdMatrixFnBatchKernel {
239    pub kind: SpdMatFn,
240}
241impl CpuKernel for SpdMatrixFnBatchKernel {
242    fn name(&self) -> &str {
243        "core.spd_matrix_fn_batch"
244    }
245    fn execute(
246        &self,
247        inputs: &[CpuTensorRef<'_>],
248        output: CpuTensorMut<'_>,
249        _attrs: &[u8],
250    ) -> Result<(), String> {
251        let x = inputs[0].expect_f64("spd_matfn X")?;
252        let (batch, n) = (dim(&inputs[0], 0), dim(&inputs[0], 1));
253        let covs: Vec<Vec<f64>> = (0..batch)
254            .map(|bi| x[bi * n * n..(bi + 1) * n * n].to_vec())
255            .collect();
256        let res = match self.kind {
257            SpdMatFn::Logm => spd::logm_batch(&covs, n),
258            SpdMatFn::Expm => spd::expm_batch(&covs, n),
259            SpdMatFn::Sqrtm => spd::sqrtm_batch(&covs, n),
260            SpdMatFn::Invsqrtm => spd::invsqrtm_batch(&covs, n),
261        };
262        let out = output.expect_f64_mut("spd_matfn Y")?;
263        out.copy_from_slice(&res.concat());
264        Ok(())
265    }
266}
267
268/// SpdLogMap backward. Inputs `[base [n,n], x [n,n], dY [n,n]]` →
269/// packed `[d_base ∥ d_x]` `[2,n,n]`.
270pub struct SpdLogMapBackwardKernel;
271impl CpuKernel for SpdLogMapBackwardKernel {
272    fn name(&self) -> &str {
273        "core.spd_log_map_backward"
274    }
275    fn execute(
276        &self,
277        inputs: &[CpuTensorRef<'_>],
278        output: CpuTensorMut<'_>,
279        _attrs: &[u8],
280    ) -> Result<(), String> {
281        let base = inputs[0].expect_f64("log_map_bwd base")?;
282        let x = inputs[1].expect_f64("log_map_bwd x")?;
283        let dy = inputs[2].expect_f64("log_map_bwd dY")?;
284        let n = dim(&inputs[0], 0);
285        let (d_base, d_x) = spd::log_map_backward(base, x, dy, n);
286        let out = output.expect_f64_mut("log_map_bwd packed")?;
287        out[..n * n].copy_from_slice(&d_base);
288        out[n * n..].copy_from_slice(&d_x);
289        Ok(())
290    }
291}
292
293/// SpdExpMap backward. Inputs `[base [n,n], v [n,n], dY [n,n]]` →
294/// packed `[d_base ∥ d_v]` `[2,n,n]`.
295pub struct SpdExpMapBackwardKernel;
296impl CpuKernel for SpdExpMapBackwardKernel {
297    fn name(&self) -> &str {
298        "core.spd_exp_map_backward"
299    }
300    fn execute(
301        &self,
302        inputs: &[CpuTensorRef<'_>],
303        output: CpuTensorMut<'_>,
304        _attrs: &[u8],
305    ) -> Result<(), String> {
306        let base = inputs[0].expect_f64("exp_map_bwd base")?;
307        let v = inputs[1].expect_f64("exp_map_bwd v")?;
308        let dy = inputs[2].expect_f64("exp_map_bwd dY")?;
309        let n = dim(&inputs[0], 0);
310        let (d_base, d_v) = spd::exp_map_backward(base, v, dy, n);
311        let out = output.expect_f64_mut("exp_map_bwd packed")?;
312        out[..n * n].copy_from_slice(&d_base);
313        out[n * n..].copy_from_slice(&d_v);
314        Ok(())
315    }
316}
317
318/// SpdParallelTransport backward. Inputs `[from, to, v, dY]` (each `[n,n]`) →
319/// packed `[d_from ∥ d_to ∥ d_v]` `[3,n,n]`.
320pub struct SpdParallelTransportBackwardKernel;
321impl CpuKernel for SpdParallelTransportBackwardKernel {
322    fn name(&self) -> &str {
323        "core.spd_parallel_transport_backward"
324    }
325    fn execute(
326        &self,
327        inputs: &[CpuTensorRef<'_>],
328        output: CpuTensorMut<'_>,
329        _attrs: &[u8],
330    ) -> Result<(), String> {
331        let from = inputs[0].expect_f64("transport_bwd from")?;
332        let to = inputs[1].expect_f64("transport_bwd to")?;
333        let v = inputs[2].expect_f64("transport_bwd v")?;
334        let dy = inputs[3].expect_f64("transport_bwd dY")?;
335        let n = dim(&inputs[0], 0);
336        let (d_from, d_to, d_v) = spd::parallel_transport_backward(from, to, v, dy, n);
337        let out = output.expect_f64_mut("transport_bwd packed")?;
338        out[..n * n].copy_from_slice(&d_from);
339        out[n * n..2 * n * n].copy_from_slice(&d_to);
340        out[2 * n * n..].copy_from_slice(&d_v);
341        Ok(())
342    }
343}
344
345/// SpdMatrixFnBatch backward. Inputs `[X [B,n,n], dY [B,n,n]]` → `dX [B,n,n]`.
346pub struct SpdMatrixFnBatchBackwardKernel {
347    pub kind: SpdMatFn,
348}
349impl CpuKernel for SpdMatrixFnBatchBackwardKernel {
350    fn name(&self) -> &str {
351        "core.spd_matrix_fn_batch_backward"
352    }
353    fn execute(
354        &self,
355        inputs: &[CpuTensorRef<'_>],
356        output: CpuTensorMut<'_>,
357        _attrs: &[u8],
358    ) -> Result<(), String> {
359        let x = inputs[0].expect_f64("spd_matfn_bwd X")?;
360        let dy = inputs[1].expect_f64("spd_matfn_bwd dY")?;
361        let n = dim(&inputs[0], 1);
362        // Same (f, f') as the batch forward (see spd::{logm,expm,sqrtm,invsqrtm}_batch).
363        let dx = match self.kind {
364            SpdMatFn::Logm => spd::matrix_fn_batch_backward(
365                x,
366                dy,
367                n,
368                |l| l.max(1e-12).ln(),
369                |l| 1.0 / l.max(1e-12),
370            ),
371            SpdMatFn::Expm => spd::matrix_fn_batch_backward(x, dy, n, |l| l.exp(), |l| l.exp()),
372            SpdMatFn::Sqrtm => spd::matrix_fn_batch_backward(
373                x,
374                dy,
375                n,
376                |l| l.max(0.0).sqrt(),
377                |l| 0.5 / l.max(1e-12).sqrt(),
378            ),
379            SpdMatFn::Invsqrtm => spd::matrix_fn_batch_backward(
380                x,
381                dy,
382                n,
383                |l| 1.0 / l.max(1e-12).sqrt(),
384                |l| -0.5 * l.max(1e-12).powf(-1.5),
385            ),
386        };
387        let out = output.expect_f64_mut("spd_matfn_bwd dX")?;
388        out.copy_from_slice(&dx);
389        Ok(())
390    }
391}
392
393/// `n` from a packed `[λ ∥ U]` length `L = n²+n` (⇒ `n = (√(1+4L)−1)/2`).
394fn n_from_packed(len: usize) -> usize {
395    (((1 + 4 * len) as f64).sqrt().round() as usize - 1) / 2
396}
397
398/// Symmetric eigendecomposition. Input `[A n,n]` → packed `[λ ∥ U]` `[n²+n]`.
399pub struct EighKernel;
400impl CpuKernel for EighKernel {
401    fn name(&self) -> &str {
402        "core.eigh"
403    }
404    fn execute(
405        &self,
406        inputs: &[CpuTensorRef<'_>],
407        output: CpuTensorMut<'_>,
408        _attrs: &[u8],
409    ) -> Result<(), String> {
410        let a = inputs[0].expect_f64("eigh A")?;
411        let n = dim(&inputs[0], 0);
412        let out = output.expect_f64_mut("eigh packed")?;
413        out.copy_from_slice(&spd::eigh_packed(a, n));
414        Ok(())
415    }
416}
417
418/// Eigh backward. Inputs `[fwd [n²+n], bar [n²+n]]` → `Ā [n,n]`.
419pub struct EighBackwardKernel;
420impl CpuKernel for EighBackwardKernel {
421    fn name(&self) -> &str {
422        "core.eigh_backward"
423    }
424    fn execute(
425        &self,
426        inputs: &[CpuTensorRef<'_>],
427        output: CpuTensorMut<'_>,
428        _attrs: &[u8],
429    ) -> Result<(), String> {
430        let fwd = inputs[0].expect_f64("eigh_bwd fwd")?;
431        let bar = inputs[1].expect_f64("eigh_bwd bar")?;
432        let n = n_from_packed(fwd.len());
433        let out = output.expect_f64_mut("eigh_bwd dA")?;
434        out.copy_from_slice(&spd::eigh_backward_packed(fwd, bar, n));
435        Ok(())
436    }
437}
438
439/// Batched eigh. Input `[A B,n,n]` → packed `[B, n²+n]`.
440pub struct EighBatchKernel;
441impl CpuKernel for EighBatchKernel {
442    fn name(&self) -> &str {
443        "core.eigh_batch"
444    }
445    fn execute(
446        &self,
447        inputs: &[CpuTensorRef<'_>],
448        output: CpuTensorMut<'_>,
449        _attrs: &[u8],
450    ) -> Result<(), String> {
451        let a = inputs[0].expect_f64("eigh_batch A")?;
452        let (b, n) = (dim(&inputs[0], 0), dim(&inputs[0], 1));
453        let covs: Vec<Vec<f64>> = (0..b)
454            .map(|bi| a[bi * n * n..(bi + 1) * n * n].to_vec())
455            .collect();
456        let out = output.expect_f64_mut("eigh_batch packed")?;
457        out.copy_from_slice(&spd::eigh_batch_packed(&covs, n).concat());
458        Ok(())
459    }
460}
461
462/// Batched eigh backward. Inputs `[fwd [B,n²+n], bar [B,n²+n]]` → `Ā [B,n,n]`.
463pub struct EighBatchBackwardKernel;
464impl CpuKernel for EighBatchBackwardKernel {
465    fn name(&self) -> &str {
466        "core.eigh_batch_backward"
467    }
468    fn execute(
469        &self,
470        inputs: &[CpuTensorRef<'_>],
471        output: CpuTensorMut<'_>,
472        _attrs: &[u8],
473    ) -> Result<(), String> {
474        use rayon::prelude::*;
475        let fwd = inputs[0].expect_f64("eigh_batch_bwd fwd")?;
476        let bar = inputs[1].expect_f64("eigh_batch_bwd bar")?;
477        let (b, packed) = (dim(&inputs[0], 0), dim(&inputs[0], 1));
478        let n = n_from_packed(packed);
479        let da: Vec<f64> = (0..b)
480            .into_par_iter()
481            .flat_map(|bi| {
482                let f = &fwd[bi * packed..(bi + 1) * packed];
483                let bb = &bar[bi * packed..(bi + 1) * packed];
484                spd::eigh_backward_packed(f, bb, n)
485            })
486            .collect();
487        let out = output.expect_f64_mut("eigh_batch_bwd dA")?;
488        out.copy_from_slice(&da);
489        Ok(())
490    }
491}
492
493/// ReEig backward. Inputs `[λ [n], U [n²], dY [n²]]` → `dX [n,n]`,
494/// reusing the forward eigendecomposition (no recompute).
495pub struct ReEigBackwardKernel {
496    pub eps: f64,
497}
498impl CpuKernel for ReEigBackwardKernel {
499    fn name(&self) -> &str {
500        "core.reeig_backward"
501    }
502    fn execute(
503        &self,
504        inputs: &[CpuTensorRef<'_>],
505        output: CpuTensorMut<'_>,
506        _attrs: &[u8],
507    ) -> Result<(), String> {
508        let lam = inputs[0].expect_f64("reeig_bwd λ")?;
509        let u = inputs[1].expect_f64("reeig_bwd U")?;
510        let dy = inputs[2].expect_f64("reeig_bwd dY")?;
511        let n = lam.len();
512        let out = output.expect_f64_mut("reeig_bwd dX")?;
513        out.copy_from_slice(&spd::reeig_backward_precomputed(lam, u, dy, n, self.eps));
514        Ok(())
515    }
516}
517
518/// LogEig backward. Inputs `[λ [n], U [n²], dY [n²]]` → `dX [n,n]`.
519pub struct LogEigBackwardKernel {
520    pub eps: f64,
521}
522impl CpuKernel for LogEigBackwardKernel {
523    fn name(&self) -> &str {
524        "core.logeig_backward"
525    }
526    fn execute(
527        &self,
528        inputs: &[CpuTensorRef<'_>],
529        output: CpuTensorMut<'_>,
530        _attrs: &[u8],
531    ) -> Result<(), String> {
532        let lam = inputs[0].expect_f64("logeig_bwd λ")?;
533        let u = inputs[1].expect_f64("logeig_bwd U")?;
534        let dy = inputs[2].expect_f64("logeig_bwd dY")?;
535        let n = lam.len();
536        let out = output.expect_f64_mut("logeig_bwd dX")?;
537        out.copy_from_slice(&spd::logeig_backward_precomputed(lam, u, dy, n, self.eps));
538        Ok(())
539    }
540}
541
542/// SPD batch-norm backward w.r.t. X. Inputs `[mean [n,n], G [n,n], dY [B,n,n]]`.
543pub struct SpdBnBackwardXKernel {
544    pub eps: f64,
545}
546impl CpuKernel for SpdBnBackwardXKernel {
547    fn name(&self) -> &str {
548        "core.spd_batch_norm_backward_x"
549    }
550    fn execute(
551        &self,
552        inputs: &[CpuTensorRef<'_>],
553        output: CpuTensorMut<'_>,
554        _attrs: &[u8],
555    ) -> Result<(), String> {
556        let mean = inputs[0].expect_f64("spd_bn_bwd_x mean")?;
557        let g = inputs[1].expect_f64("spd_bn_bwd_x G")?;
558        let dy = inputs[2].expect_f64("spd_bn_bwd_x dY")?;
559        let (batch, n) = (dim(&inputs[2], 0), dim(&inputs[0], 0));
560        let out = output.expect_f64_mut("spd_bn_bwd_x dX")?;
561        out.copy_from_slice(&spd::spd_bn_backward_x(mean, g, dy, batch, n, self.eps));
562        Ok(())
563    }
564}
565
566/// SPD batch-norm backward w.r.t. G. Inputs `[X [B,n,n], mean [n,n], G [n,n], dY [B,n,n]]`.
567pub struct SpdBnBackwardGKernel {
568    pub eps: f64,
569}
570impl CpuKernel for SpdBnBackwardGKernel {
571    fn name(&self) -> &str {
572        "core.spd_batch_norm_backward_g"
573    }
574    fn execute(
575        &self,
576        inputs: &[CpuTensorRef<'_>],
577        output: CpuTensorMut<'_>,
578        _attrs: &[u8],
579    ) -> Result<(), String> {
580        let x = inputs[0].expect_f64("spd_bn_bwd_g X")?;
581        let mean = inputs[1].expect_f64("spd_bn_bwd_g mean")?;
582        let g = inputs[2].expect_f64("spd_bn_bwd_g G")?;
583        let dy = inputs[3].expect_f64("spd_bn_bwd_g dY")?;
584        let (batch, n) = (dim(&inputs[0], 0), dim(&inputs[0], 1));
585        let out = output.expect_f64_mut("spd_bn_bwd_g dG")?;
586        out.copy_from_slice(&spd::spd_bn_backward_g(x, mean, g, dy, batch, n, self.eps));
587        Ok(())
588    }
589}
rlx_cpu/spd_kernels.rs

rlx_cpu/
spd_kernels.rs
