1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
//! Fréchet sensitivity w.r.t. individual CSR entries of a symmetric operator (ADR-0186, §55.5).
//!
//! [`EntrySensitivity`] implements [`GeneratorSensitivity`] for a stencil `∂A/∂L_{ij}`.
//! Plugs into the unchanged `graph_expmv_frechet` infrastructure (A2).
use alloc::vec::Vec;
use crate::{
error::SemiflowError,
float::SemiflowFloat,
graph_sensitivity::{check_bilinear_args, GeneratorSensitivity},
scratch::ScratchPool,
};
/// Parameter list for single-entry symmetric-stencil sensitivity (§55.5).
///
/// Each entry `(i, j)` in `entries` (with `i ≤ j`) represents the pair `{L_{ij}, L_{ji}}`
/// which are tied by symmetry. `n_nodes` is the operator dimension.
pub struct EntrySensitivity {
/// List of `(row, col)` pairs with `row ≤ col`. One parameter per entry.
pub entries: Vec<(usize, usize)>,
/// Operator dimension (number of nodes).
pub n_nodes: usize,
}
impl<F: SemiflowFloat> GeneratorSensitivity<F> for EntrySensitivity {
fn n_params(&self) -> usize {
self.entries.len()
}
/// `out ← (∂A/∂L_{ij}) · v` for the k-th parameter.
///
/// Since `A = −L` (ADR-0186 sign convention), `∂A/∂L_{ij} = −∂L/∂L_{ij}`.
///
/// Stencil (§55.5):
/// - Diagonal (`i == j`): `(∂L/∂L_{ii}) v = v[i] e_i`, so `out[i] = −v[i]`.
/// - Off-diagonal (`i ≠ j`): `(∂L/∂L_{ij}) v = v[j] e_i + v[i] e_j`,
/// so `out[i] = −v[j]`, `out[j] = −v[i]`.
///
/// # Errors
///
/// [`SemiflowError::DomainViolation`] if `k ≥ self.n_params()`.
fn apply_param_deriv(
&self,
k: usize,
_t: F,
v: &[F],
out: &mut [F],
) -> Result<(), SemiflowError> {
if k >= self.entries.len() {
return Err(SemiflowError::DomainViolation {
what: "EntrySensitivity::apply_param_deriv: k out of range",
#[allow(clippy::cast_precision_loss)]
value: k as f64,
});
}
for x in out.iter_mut() {
*x = F::zero();
}
let (i, j) = self.entries[k];
if i == j {
out[i] = -v[i]; // diagonal entry: ∂A/∂L_{ii} = −e_i eᵢᵀ
} else {
out[i] = -v[j]; // off-diagonal: ∂A/∂L_{ij} = −(e_i eⱼᵀ + e_j eᵢᵀ)
out[j] = -v[i];
}
Ok(())
}
/// `O(n_params)` contraction (§63.6.b): off-diagonal
/// `grad[k] −= w·(a_i b_j + a_j b_i)`, diagonal `grad[k] −= w·a_i b_i`.
fn accumulate_bilinear(
&self,
_t: F,
w: F,
a: &[F],
b: &[F],
grad: &mut [F],
_scratch: &mut ScratchPool<F>,
) -> Result<(), SemiflowError> {
check_bilinear_args(a.len(), b.len(), grad.len(), self.entries.len())?;
let n = a.len();
for (gk, &(i, j)) in grad.iter_mut().zip(&self.entries) {
if i >= n || j >= n {
return Err(SemiflowError::DomainViolation {
what: "EntrySensitivity::accumulate_bilinear: index out of range",
#[allow(clippy::cast_precision_loss)]
value: i.max(j) as f64,
});
}
*gk -= if i == j {
w * a[i] * b[i]
} else {
w * (a[i] * b[j] + a[j] * b[i])
};
}
Ok(())
}
}