pub struct SparseTensor {
pub values: NodeId,
pub col_idx: NodeId,
pub row_ptr: NodeId,
pub n_rows: usize,
pub n_cols: usize,
}Expand description
CSR-format sparse matrix at the IR level. Bundles the three
CSR NodeIds with structural shape info known at graph-build time.
Modeled on jax.experimental.sparse.BCOO — a wrapper around
(data, indices) arrays plus a shape tuple, with methods that
expand to the right subgraph for each operation.
Fields§
§values: NodeIdNon-zero values in row-major CSR order. F64.
col_idx: NodeIdColumn index per non-zero. I32.
row_ptr: NodeIdStart index in values / col_idx per row, length n_rows + 1. I32.
n_rows: usizeLogical row count of A.
n_cols: usizeLogical column count of A (n_rows == n_cols for square / SPD).
Implementations§
Source§impl SparseTensor
impl SparseTensor
Sourcepub fn from_csr(
values: NodeId,
col_idx: NodeId,
row_ptr: NodeId,
n_rows: usize,
n_cols: usize,
) -> Self
pub fn from_csr( values: NodeId, col_idx: NodeId, row_ptr: NodeId, n_rows: usize, n_cols: usize, ) -> Self
Build from existing CSR NodeIds. Caller is responsible for
the layout invariants (sortedness within rows,
row_ptr.len() == n_rows + 1, etc.).
Sourcepub fn mat_vec(&self, g: &mut Graph, x: NodeId) -> NodeId
pub fn mat_vec(&self, g: &mut Graph, x: NodeId) -> NodeId
y = A · x for a length-n_cols dense vector.
Sourcepub fn cg_solve(
&self,
g: &mut Graph,
b: NodeId,
max_iter: u32,
tol: f64,
) -> NodeId
pub fn cg_solve( &self, g: &mut Graph, b: NodeId, max_iter: u32, tol: f64, ) -> NodeId
x = A⁻¹ · b via Conjugate Gradient. SPD only. tol is the
absolute residual threshold; max_iter caps iteration count.
Sourcepub fn solve_general(
&self,
g: &mut Graph,
b: NodeId,
adjoint: &SparseTensor,
) -> NodeId
pub fn solve_general( &self, g: &mut Graph, b: NodeId, adjoint: &SparseTensor, ) -> NodeId
x = A⁻¹ · b via direct LU for non-symmetric A. The
caller supplies an explicit transpose adjoint (CSR of Aᵀ)
— for square non-symmetric matrices, Aᵀ has the same
nnz pattern as A after a CSR↔CSC swap. The forward solve
uses only self; the VJP routes the adjoint solve through
adjoint. Use this in place of solve when the assumption
Aᵀ = A would be wrong.
Sourcepub fn pcg_solve(
&self,
g: &mut Graph,
b: NodeId,
max_iter: u32,
tol: f64,
) -> NodeId
pub fn pcg_solve( &self, g: &mut Graph, b: NodeId, max_iter: u32, tol: f64, ) -> NodeId
x = A⁻¹ · b via Jacobi-preconditioned CG. SPD only.
Convergence dramatically faster than plain CG on ill-
conditioned matrices where diag(A) captures most of the
magnitude variation — typical for circuit MNA matrices with
mixed-magnitude device parameters. The preconditioner is
extracted from the CSR by the kernel; no separate input.
Sourcepub fn transpose_values(
&self,
g: &mut Graph,
col_idx_t: NodeId,
row_ptr_t: NodeId,
) -> NodeId
pub fn transpose_values( &self, g: &mut Graph, col_idx_t: NodeId, row_ptr_t: NodeId, ) -> NodeId
Permute this tensor’s values into the values vector of
Aᵀ. The transposed pattern (col_idx_t, row_ptr_t) is
supplied as NodeIds — typically computed once via
crate::csr_transpose_pattern and embedded as Op::Constant
since the pattern is fixed across Newton iterations.
Sourcepub fn cholesky_solve(&self, g: &mut Graph, b: NodeId) -> NodeId
pub fn cholesky_solve(&self, g: &mut Graph, b: NodeId) -> NodeId
x = A⁻¹ · b via direct sparse Cholesky for SPD A. Densifies
into a dense buffer and calls LAPACK dpotrf + triangular
solves. Mirror of solve (LU-based) but ½× factor cost and
numerically more stable; only valid when A is SPD.
Sourcepub fn lsqr_solve(
&self,
g: &mut Graph,
b: NodeId,
max_iter: u32,
tol: f64,
) -> NodeId
pub fn lsqr_solve( &self, g: &mut Graph, b: NodeId, max_iter: u32, tol: f64, ) -> NodeId
x = argmin ||A·x - b||₂ via LSQR (Paige-Saunders 1982).
Works for any A (square / over-determined / under-determined);
returns the minimum-norm solution when A is rank-deficient or
under-determined. VJP not implemented in v1.
Sourcepub fn bicgstab_solve(
&self,
g: &mut Graph,
b: NodeId,
max_iter: u32,
tol: f64,
) -> NodeId
pub fn bicgstab_solve( &self, g: &mut Graph, b: NodeId, max_iter: u32, tol: f64, ) -> NodeId
x = A⁻¹ · b via BiCGSTAB for non-symmetric A. Same
shape contract as CG/PCG (no explicit adjoint pattern needed —
the kernel itself can solve Aᵀ·x = b via a flag, used by VJPs).
Sourcepub fn ilu_pcg_solve(
&self,
g: &mut Graph,
b: NodeId,
max_iter: u32,
tol: f64,
) -> NodeId
pub fn ilu_pcg_solve( &self, g: &mut Graph, b: NodeId, max_iter: u32, tol: f64, ) -> NodeId
x = A⁻¹ · b via ILU(0)-preconditioned CG. SPD A required
(same contract as CG/PCG). ILU is factored on each call —
for static-pattern Newton loops the cost amortizes against the
faster convergence vs. Jacobi-PCG.
Sourcepub fn gmres_solve(
&self,
g: &mut Graph,
b: NodeId,
max_iter: u32,
tol: f64,
adjoint: &SparseTensor,
) -> NodeId
pub fn gmres_solve( &self, g: &mut Graph, b: NodeId, max_iter: u32, tol: f64, adjoint: &SparseTensor, ) -> NodeId
x = A⁻¹ · b via GMRES for non-symmetric A. Same
transpose-triplet contract as solve_general. max_iter
caps Krylov dimension; tol is the residual norm threshold.
Trait Implementations§
Source§impl Clone for SparseTensor
impl Clone for SparseTensor
Source§fn clone(&self) -> SparseTensor
fn clone(&self) -> SparseTensor
1.0.0 (const: unstable) · Source§fn clone_from(&mut self, source: &Self)
fn clone_from(&mut self, source: &Self)
source. Read moreimpl Copy for SparseTensor
Auto Trait Implementations§
impl Freeze for SparseTensor
impl RefUnwindSafe for SparseTensor
impl Send for SparseTensor
impl Sync for SparseTensor
impl Unpin for SparseTensor
impl UnsafeUnpin for SparseTensor
impl UnwindSafe for SparseTensor
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Source§impl<T> CloneToUninit for Twhere
T: Clone,
impl<T> CloneToUninit for Twhere
T: Clone,
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self>
fn into_either(self, into_left: bool) -> Either<Self, Self>
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more