rlx_ir/ops/
spd_graph.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! **All-backend SPD-manifold layers via the graph-Jacobi eigensolver.**
17//!
18//! The native SPD ops `Op::EighBatch` / `Op::SpdMatrixFnBatch` / `Op::SpdLogMap`
19//! (see [`super::manifold`]) currently only carry Metal + CUDA kernels, so a
20//! model that needs the AIRM manifold on CoreML / MLX cannot use them. This
21//! module expresses the same operations purely through the backend-agnostic
22//! graph eigensolver in [`super::spd_eig`] (cyclic Jacobi via `Op::Scan` +
23//! `mm`/`sqrt`/`Activation::Log`), so they lower to **every** backend — the same
24//! path the `rlx-repspd` / `rlx-mpnet` SPD-manifold ports use to reach all five
25//! Mac backends (+ CUDA).
26//!
27//! These are additive `*_graph` variants of the manifold builders; the native
28//! ops remain for CPU-f64-LAPACK / GPU-kernel fast paths where they exist.
29
30use crate::infer::GraphExt;
31use crate::op::SpdMatFn;
32use crate::ops::spd_eig::{SpectralFn, spd_jacobi_sweeps, spectral_matfn_batched};
33use crate::{Graph, NodeId};
34
35impl Graph {
36    /// Batched SPD spectral matrix function via the **all-backend** graph
37    /// eigensolver — the portable counterpart of [`Graph::spd_matrix_fn_batch`]
38    /// (which emits the Metal/CUDA-only `Op::SpdMatrixFnBatch`).
39    ///
40    /// `x` is `[batch, n, n]` symmetric-PSD; applies `kind` (logm / sqrtm /
41    /// invsqrtm) to each matrix and returns `[batch, n, n]`. `SpdMatFn::Expm`
42    /// has no Jacobi-graph form here and falls back to the native op.
43    ///
44    /// `eps` floors the eigenvalues (`max(λ, eps)`) exactly as
45    /// [`super::spd_eig::spectral_matfn_batched`].
46    pub fn spd_matrix_fn_batch_graph(&mut self, x: NodeId, kind: SpdMatFn, eps: f64) -> NodeId {
47        let xs = self.node(x).shape.clone();
48        let b = xs.dim(0).unwrap_static();
49        let n = xs.dim(1).unwrap_static();
50        let sweeps = spd_jacobi_sweeps();
51        let f = match kind {
52            SpdMatFn::Logm => SpectralFn::Log,
53            SpdMatFn::Sqrtm => SpectralFn::Sqrt,
54            SpdMatFn::Invsqrtm => SpectralFn::InvSqrt,
55            // Expm has no Jacobi-graph spectral variant here — use the native op.
56            SpdMatFn::Expm => return self.spd_matrix_fn_batch(x, SpdMatFn::Expm),
57        };
58        spectral_matfn_batched(self, x, b, n, sweeps, eps, f)
59    }
60
61    /// Batched matrix logarithm via the all-backend graph eigensolver.
62    pub fn spd_logm_batch_graph(&mut self, x: NodeId, eps: f64) -> NodeId {
63        self.spd_matrix_fn_batch_graph(x, SpdMatFn::Logm, eps)
64    }
65    /// Batched matrix square root via the all-backend graph eigensolver.
66    pub fn spd_sqrtm_batch_graph(&mut self, x: NodeId, eps: f64) -> NodeId {
67        self.spd_matrix_fn_batch_graph(x, SpdMatFn::Sqrtm, eps)
68    }
69    /// Batched inverse matrix square root via the all-backend graph eigensolver.
70    pub fn spd_invsqrtm_batch_graph(&mut self, x: NodeId, eps: f64) -> NodeId {
71        self.spd_matrix_fn_batch_graph(x, SpdMatFn::Invsqrtm, eps)
72    }
73
74    /// AIRM Riemannian logarithm at an arbitrary base, **all backends**:
75    /// `Log_P(X) = P^{1/2} · logm(P^{-1/2} X P^{-1/2}) · P^{1/2}`.
76    ///
77    /// The portable counterpart of [`Graph::spd_log_map`] (native `Op::SpdLogMap`,
78    /// Metal/CUDA-only). Both `base` and `x` are `[batch, n, n]` SPD; output is
79    /// the `[batch, n, n]` tangent vector at `base`. This is the AIRM op the
80    /// Log-Euclidean / RepSPD-style manifold cross-attention heads need on
81    /// CoreML / MLX.
82    pub fn spd_log_map_graph(&mut self, base: NodeId, x: NodeId, eps: f64) -> NodeId {
83        let p_half = self.spd_sqrtm_batch_graph(base, eps); // P^{1/2}
84        let p_ihalf = self.spd_invsqrtm_batch_graph(base, eps); // P^{-1/2}
85        // inner = P^{-1/2} · X · P^{-1/2}   (batched mm)
86        let px = self.mm(p_ihalf, x);
87        let inner = self.mm(px, p_ihalf);
88        let log_inner = self.spd_logm_batch_graph(inner, eps);
89        // Log_P(X) = P^{1/2} · log_inner · P^{1/2}
90        let ph_log = self.mm(p_half, log_inner);
91        self.mm(ph_log, p_half)
92    }
93
94    /// AIRM Riemannian exponential at an arbitrary base, **all backends**:
95    /// `Exp_P(V) = P^{1/2} · expm(P^{-1/2} V P^{-1/2}) · P^{1/2}` — inverse of
96    /// [`Graph::spd_log_map_graph`]. Uses `P^{1/2}`/`P^{-1/2}` from the graph
97    /// eigensolver; the inner `expm` uses the native `Op::SpdMatrixFnBatch`
98    /// (`Expm` has no Jacobi-graph spectral form), so this is all-backend only
99    /// where `expm` is available — prefer the log map for portability.
100    pub fn spd_exp_map_graph(&mut self, base: NodeId, v: NodeId, eps: f64) -> NodeId {
101        let p_half = self.spd_sqrtm_batch_graph(base, eps);
102        let p_ihalf = self.spd_invsqrtm_batch_graph(base, eps);
103        let pv = self.mm(p_ihalf, v);
104        let inner = self.mm(pv, p_ihalf);
105        let exp_inner = self.spd_matrix_fn_batch_graph(inner, SpdMatFn::Expm, eps);
106        let ph_exp = self.mm(p_half, exp_inner);
107        self.mm(ph_exp, p_half)
108    }
109}
rlx_ir/ops/spd_graph.rs

rlx_ir/ops/
spd_graph.rs
