rlx_ir/ops/
spd_graph.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//! **All-backend SPD-manifold layers via the graph-Jacobi eigensolver.**
6//!
7//! The native SPD ops `Op::EighBatch` / `Op::SpdMatrixFnBatch` / `Op::SpdLogMap`
8//! (see [`super::manifold`]) currently only carry Metal + CUDA kernels, so a
9//! model that needs the AIRM manifold on CoreML / MLX cannot use them. This
10//! module expresses the same operations purely through the backend-agnostic
11//! graph eigensolver in [`super::spd_eig`] (cyclic Jacobi via `Op::Scan` +
12//! `mm`/`sqrt`/`Activation::Log`), so they lower to **every** backend — the same
13//! path the `rlx-repspd` / `rlx-mpnet` SPD-manifold ports use to reach all five
14//! Mac backends (+ CUDA).
15//!
16//! These are additive `*_graph` variants of the manifold builders; the native
17//! ops remain for CPU-f64-LAPACK / GPU-kernel fast paths where they exist.
18
19use crate::infer::GraphExt;
20use crate::op::SpdMatFn;
21use crate::ops::spd_eig::{SpectralFn, spd_jacobi_sweeps, spectral_matfn_batched};
22use crate::{Graph, NodeId};
23
24impl Graph {
25    /// Batched SPD spectral matrix function via the **all-backend** graph
26    /// eigensolver — the portable counterpart of [`Graph::spd_matrix_fn_batch`]
27    /// (which emits the Metal/CUDA-only `Op::SpdMatrixFnBatch`).
28    ///
29    /// `x` is `[batch, n, n]` symmetric-PSD; applies `kind` (logm / sqrtm /
30    /// invsqrtm) to each matrix and returns `[batch, n, n]`. `SpdMatFn::Expm`
31    /// has no Jacobi-graph form here and falls back to the native op.
32    ///
33    /// `eps` floors the eigenvalues (`max(λ, eps)`) exactly as
34    /// [`super::spd_eig::spectral_matfn_batched`].
35    pub fn spd_matrix_fn_batch_graph(&mut self, x: NodeId, kind: SpdMatFn, eps: f64) -> NodeId {
36        let xs = self.node(x).shape.clone();
37        let b = xs.dim(0).unwrap_static();
38        let n = xs.dim(1).unwrap_static();
39        let sweeps = spd_jacobi_sweeps();
40        let f = match kind {
41            SpdMatFn::Logm => SpectralFn::Log,
42            SpdMatFn::Sqrtm => SpectralFn::Sqrt,
43            SpdMatFn::Invsqrtm => SpectralFn::InvSqrt,
44            // Expm has no Jacobi-graph spectral variant here — use the native op.
45            SpdMatFn::Expm => return self.spd_matrix_fn_batch(x, SpdMatFn::Expm),
46        };
47        spectral_matfn_batched(self, x, b, n, sweeps, eps, f)
48    }
49
50    /// Batched matrix logarithm via the all-backend graph eigensolver.
51    pub fn spd_logm_batch_graph(&mut self, x: NodeId, eps: f64) -> NodeId {
52        self.spd_matrix_fn_batch_graph(x, SpdMatFn::Logm, eps)
53    }
54    /// Batched matrix square root via the all-backend graph eigensolver.
55    pub fn spd_sqrtm_batch_graph(&mut self, x: NodeId, eps: f64) -> NodeId {
56        self.spd_matrix_fn_batch_graph(x, SpdMatFn::Sqrtm, eps)
57    }
58    /// Batched inverse matrix square root via the all-backend graph eigensolver.
59    pub fn spd_invsqrtm_batch_graph(&mut self, x: NodeId, eps: f64) -> NodeId {
60        self.spd_matrix_fn_batch_graph(x, SpdMatFn::Invsqrtm, eps)
61    }
62
63    /// AIRM Riemannian logarithm at an arbitrary base, **all backends**:
64    /// `Log_P(X) = P^{1/2} · logm(P^{-1/2} X P^{-1/2}) · P^{1/2}`.
65    ///
66    /// The portable counterpart of [`Graph::spd_log_map`] (native `Op::SpdLogMap`,
67    /// Metal/CUDA-only). Both `base` and `x` are `[batch, n, n]` SPD; output is
68    /// the `[batch, n, n]` tangent vector at `base`. This is the AIRM op the
69    /// Log-Euclidean / RepSPD-style manifold cross-attention heads need on
70    /// CoreML / MLX.
71    pub fn spd_log_map_graph(&mut self, base: NodeId, x: NodeId, eps: f64) -> NodeId {
72        let p_half = self.spd_sqrtm_batch_graph(base, eps); // P^{1/2}
73        let p_ihalf = self.spd_invsqrtm_batch_graph(base, eps); // P^{-1/2}
74        // inner = P^{-1/2} · X · P^{-1/2}   (batched mm)
75        let px = self.mm(p_ihalf, x);
76        let inner = self.mm(px, p_ihalf);
77        let log_inner = self.spd_logm_batch_graph(inner, eps);
78        // Log_P(X) = P^{1/2} · log_inner · P^{1/2}
79        let ph_log = self.mm(p_half, log_inner);
80        self.mm(ph_log, p_half)
81    }
82
83    /// AIRM Riemannian exponential at an arbitrary base, **all backends**:
84    /// `Exp_P(V) = P^{1/2} · expm(P^{-1/2} V P^{-1/2}) · P^{1/2}` — inverse of
85    /// [`Graph::spd_log_map_graph`]. Uses `P^{1/2}`/`P^{-1/2}` from the graph
86    /// eigensolver; the inner `expm` uses the native `Op::SpdMatrixFnBatch`
87    /// (`Expm` has no Jacobi-graph spectral form), so this is all-backend only
88    /// where `expm` is available — prefer the log map for portability.
89    pub fn spd_exp_map_graph(&mut self, base: NodeId, v: NodeId, eps: f64) -> NodeId {
90        let p_half = self.spd_sqrtm_batch_graph(base, eps);
91        let p_ihalf = self.spd_invsqrtm_batch_graph(base, eps);
92        let pv = self.mm(p_ihalf, v);
93        let inner = self.mm(pv, p_ihalf);
94        let exp_inner = self.spd_matrix_fn_batch_graph(inner, SpdMatFn::Expm, eps);
95        let ph_exp = self.mm(p_half, exp_inner);
96        self.mm(ph_exp, p_half)
97    }
98}
rlx_ir/ops/spd_graph.rs

rlx_ir/ops/
spd_graph.rs
