compreSSM 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,116 @@
1
+ Metadata-Version: 2.4
2
+ Name: compreSSM
3
+ Version: 1.0.0
4
+ Summary: Library built off of the compreSSM paper linked in the README
5
+ Author: Christian Rafferty
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/cfr081709/compreSSM_library
8
+ Project-URL: Repository, https://github.com/cfr081709/compreSSM_library
9
+ Project-URL: Issues, https://github.com/cfr081709/compreSSM_library/issues
10
+ Requires-Python: >=3.10
11
+ Description-Content-Type: text/markdown
12
+ License-File: LICENSE
13
+ Requires-Dist: numpy>=1.24
14
+ Requires-Dist: pydantic>=2.0
15
+ Requires-Dist: scipy>=1.10
16
+ Requires-Dist: torch>=2.0
17
+ Provides-Extra: test
18
+ Requires-Dist: pytest>=7.0; extra == "test"
19
+ Dynamic: license-file
20
+
21
+ # CompreSSM Library
22
+
23
+ CompreSSM provides balanced-truncation tools for compressing stable discrete-time
24
+ state-space models, including diagonal systems used by Linear Recurrent Units
25
+ (LRUs). It computes controllability and observability Gramians, Hankel singular
26
+ values, reduced systems, and an error bound for each reduction.
27
+
28
+ ## Installation
29
+
30
+ ```bash
31
+ python -m pip install compreSSM-library
32
+ ```
33
+
34
+ The package requires Python 3.10 or later. PyTorch is a runtime dependency
35
+ because the module also provides helpers for extracting and updating PyTorch
36
+ modules.
37
+
38
+ ## Quick start
39
+
40
+ The published module is named `core`:
41
+
42
+ ```python
43
+ import numpy as np
44
+
45
+ from core import (
46
+ CompressionConfig,
47
+ DiagonalDiscreteSSM,
48
+ compress_diagonal_numpy,
49
+ )
50
+
51
+ system = DiagonalDiscreteSSM.from_numpy(
52
+ lam=np.array([0.5, 0.8], dtype=complex),
53
+ B=np.array([[1.0], [0.01]], dtype=complex),
54
+ C=np.array([[1.0, 0.01]], dtype=complex),
55
+ D=np.array([0.0]),
56
+ )
57
+
58
+ reduced = compress_diagonal_numpy(
59
+ system,
60
+ config=CompressionConfig(tol=0.05, max_ratio=0.95),
61
+ )
62
+
63
+ if reduced is None:
64
+ print("The requested tolerance does not justify reducing this system.")
65
+ else:
66
+ print(f"Reduced state dimension: {reduced.r} (from {system.n})")
67
+ print(f"Balanced-truncation error bound: {reduced.error_bound}")
68
+ ```
69
+
70
+ `tol` is the fraction of total Hankel singular value mass that may be discarded.
71
+ `max_ratio` prevents accepting a reduction when the resulting state count is not
72
+ small enough compared with the original system.
73
+
74
+ ## Main API
75
+
76
+ - `DiscreteSSM` and `compress_numpy` handle general stable discrete-time systems.
77
+ - `DiagonalDiscreteSSM` and `compress_diagonal_numpy` handle diagonal systems,
78
+ including LRU parameterizations.
79
+ - `CompressionConfig` controls the sample time, retained Hankel singular value
80
+ mass, and minimum worthwhile reduction.
81
+ - `compress_module` and `compress_lru_module` adapt compatible PyTorch modules.
82
+ - `run_scheduled_reduction` applies reduction to blocks on a training schedule.
83
+
84
+ The `D` (direct-feedthrough) matrix is retained unchanged during reduction. The
85
+ reported balanced-truncation error bound assumes a stable, minimal system.
86
+
87
+ For details on every class and function, see the [full API documentation](docs.md).
88
+
89
+ ## Development
90
+
91
+ Install the package and test dependencies, then run the test suite:
92
+
93
+ ```bash
94
+ python -m pip install ".[test]"
95
+ python -m pytest
96
+ ```
97
+
98
+ ## Citation
99
+
100
+ This library is based on:
101
+
102
+ > Makram Chahine, Philipp Nazari, Daniela Rus, and T. Konstantin Rusch.
103
+ > “The Curious Case of In-Training Compression of State Space Models.” 2026.
104
+ > [arXiv:2510.02823](https://arxiv.org/abs/2510.02823).
105
+
106
+ ```bibtex
107
+ @misc{chahine2026curiouscaseintrainingcompression,
108
+ title={The Curious Case of In-Training Compression of State Space Models},
109
+ author={Makram Chahine and Philipp Nazari and Daniela Rus and T. Konstantin Rusch},
110
+ year={2026},
111
+ eprint={2510.02823},
112
+ archivePrefix={arXiv},
113
+ primaryClass={cs.LG},
114
+ url={https://arxiv.org/abs/2510.02823}
115
+ }
116
+ ```
@@ -0,0 +1,6 @@
1
+ core.py,sha256=1gse_gARzt9wgaEOqpDCXKukhXdQCzv8e3miMiSnafM,15127
2
+ compressm-1.0.0.dist-info/licenses/LICENSE,sha256=gYmCMmN-UsbtPaCxDWbUgpFREb5ydCJtE7bwJFO373o,1065
3
+ compressm-1.0.0.dist-info/METADATA,sha256=rD9kixU_nWwqqKpdleOtQsghW4r279t64l6reH0_tQI,3672
4
+ compressm-1.0.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
5
+ compressm-1.0.0.dist-info/top_level.txt,sha256=0ebp8k35dOTHVwbY3UcokRsjC7cWLN_iaJ70tIcmODU,5
6
+ compressm-1.0.0.dist-info/RECORD,,
@@ -0,0 +1,5 @@
1
+ Wheel-Version: 1.0
2
+ Generator: setuptools (84.2.1)
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any
5
+
@@ -0,0 +1,7 @@
1
+ Copyright 2026 Christian Rafferty
2
+
3
+ Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
4
+
5
+ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
6
+
7
+ THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
@@ -0,0 +1 @@
1
+ core
core.py ADDED
@@ -0,0 +1,396 @@
1
+ import numpy as np
2
+ import torch
3
+ from collections.abc import Callable, Iterable
4
+ from dataclasses import dataclass
5
+ from torch import nn
6
+ from pydantic import BaseModel, Field
7
+ from scipy.linalg import solve_discrete_lyapunov, cholesky, svd
8
+
9
+
10
+ class CompressionConfig(BaseModel):
11
+ dt: float = Field(default=1.0, gt=0.0)
12
+ tol: float = Field(
13
+ default=0.01,
14
+ ge=0.0,
15
+ le=1.0,
16
+ description="Fraction of the sum of Hankel singular values to discard",
17
+ )
18
+ max_ratio: float = Field(default=0.95, gt=0.0, le=1.0)
19
+
20
+ @property
21
+ def threshold(self) -> float:
22
+ """Fraction of the sum of Hankel singular values to retain."""
23
+ return 1.0 - self.tol
24
+
25
+
26
+ class ReductionSchedule(BaseModel):
27
+ red_start: int = Field(ge=0)
28
+ red_end: int = Field(ge=0)
29
+ red_interval: int = Field(gt=0)
30
+
31
+ def fires_at(self, step: int) -> bool:
32
+ return (
33
+ self.red_start <= step < self.red_end
34
+ and (step + 1) % self.red_interval == 0
35
+ )
36
+
37
+
38
+ @dataclass(frozen=True)
39
+ class DiscreteSSM:
40
+ A: np.ndarray
41
+ B: np.ndarray
42
+ C: np.ndarray
43
+ D: np.ndarray
44
+
45
+ @property
46
+ def n(self) -> int:
47
+ return self.A.shape[0]
48
+
49
+ @classmethod
50
+ def from_numpy(cls, *, A: np.ndarray, B: np.ndarray, C: np.ndarray, D: np.ndarray):
51
+ A, B, C, D = (np.asarray(matrix) for matrix in (A, B, C, D))
52
+ if A.ndim != 2 or A.shape[0] != A.shape[1] or A.shape[0] == 0:
53
+ raise ValueError("A must be a non-empty square matrix")
54
+ if B.ndim != 2 or B.shape[0] != A.shape[0]:
55
+ raise ValueError("B must be a matrix with the same number of rows as A")
56
+ if C.ndim != 2 or C.shape[1] != A.shape[0]:
57
+ raise ValueError("C must be a matrix with the same number of columns as A")
58
+ if not all(np.all(np.isfinite(matrix)) for matrix in (A, B, C, D)):
59
+ raise ValueError("State-space matrices must contain only finite values")
60
+ if np.max(np.abs(np.linalg.eigvals(A))) >= 1:
61
+ raise ValueError("A is not stable (spectral radius >= 1)")
62
+ return cls(A=A, B=B, C=C, D=D)
63
+
64
+ @classmethod
65
+ def from_continuous_diagonal_numpy(cls, *, a: np.ndarray, B: np.ndarray, C: np.ndarray, D: np.ndarray, dt: float):
66
+ """Discretize a stable diagonal continuous-time system with ZOH input."""
67
+ a = np.asarray(a)
68
+ if a.ndim != 1 or a.size == 0:
69
+ raise ValueError("a must be a non-empty one-dimensional array")
70
+ if not np.isfinite(dt) or dt <= 0:
71
+ raise ValueError("dt must be finite and greater than zero")
72
+ if not np.all(np.isfinite(a)) or np.any(np.real(a) >= 0):
73
+ raise ValueError("Continuous-time diagonal rates must be finite and stable")
74
+
75
+ a_bar = np.exp(dt * a)
76
+ integration_factor = np.full(a.shape, dt, dtype=np.result_type(a, dt))
77
+ np.divide(np.expm1(dt * a), a, out=integration_factor, where=a != 0)
78
+ B_bar = integration_factor[:, None] * B
79
+ return cls.from_numpy(A=np.diag(a_bar), B=B_bar, C=C, D=D)
80
+
81
+ @classmethod
82
+ def from_tensors(cls, *, A_log: torch.Tensor, B: torch.Tensor, C: torch.Tensor, D: torch.Tensor, dt: float):
83
+ a = -np.exp(tensor_to_numpy(A_log))
84
+ n = a.shape[0]
85
+ return cls.from_continuous_diagonal_numpy(
86
+ a=a,
87
+ B=tensor_to_numpy(B).reshape(n, -1),
88
+ C=tensor_to_numpy(C).reshape(-1, n),
89
+ D=tensor_to_numpy(D),
90
+ dt=dt,
91
+ )
92
+
93
+ @classmethod
94
+ def from_module(cls, module: nn.Module, *, dt: float):
95
+ sd = module.state_dict()
96
+ return cls.from_tensors(
97
+ A_log=sd['ssm.A_log'], B=sd['ssm.B'], C=sd['ssm.C'], D=sd['ssm.D'], dt=dt
98
+ )
99
+
100
+
101
+ @dataclass(frozen=True)
102
+ class DiagonalDiscreteSSM:
103
+ lam: np.ndarray
104
+ B: np.ndarray
105
+ C: np.ndarray
106
+ D: np.ndarray
107
+
108
+ @property
109
+ def n(self) -> int:
110
+ return self.lam.shape[0]
111
+
112
+ @classmethod
113
+ def from_numpy(cls, *, lam: np.ndarray, B: np.ndarray, C: np.ndarray, D: np.ndarray):
114
+ lam, B, C, D = (np.asarray(matrix) for matrix in (lam, B, C, D))
115
+ if lam.ndim != 1 or lam.size == 0:
116
+ raise ValueError("lam must be a non-empty one-dimensional array")
117
+ if B.ndim != 2 or B.shape[0] != lam.shape[0]:
118
+ raise ValueError("B must be a matrix with one row per state")
119
+ if C.ndim != 2 or C.shape[1] != lam.shape[0]:
120
+ raise ValueError("C must be a matrix with one column per state")
121
+ if not all(np.all(np.isfinite(matrix)) for matrix in (lam, B, C, D)):
122
+ raise ValueError("State-space matrices must contain only finite values")
123
+ if np.max(np.abs(lam)) >= 1:
124
+ raise ValueError("lam is not stable (spectral radius >= 1)")
125
+ return cls(lam=lam, B=B, C=C, D=D)
126
+
127
+ @classmethod
128
+ def from_lru_numpy(cls, *, nu_log, theta_log, B_re, B_im, gamma_log, C_re, C_im, D):
129
+ lam = np.exp(-np.exp(nu_log) + 1j * np.exp(theta_log))
130
+ B = (B_re + 1j * B_im) * np.exp(gamma_log)[:, None]
131
+ C = C_re + 1j * C_im
132
+ return cls.from_numpy(lam=lam, B=B, C=C, D=D)
133
+
134
+ @classmethod
135
+ def from_lru_module(cls, module: nn.Module):
136
+ p = {k: tensor_to_numpy(v) for k, v in module.named_parameters()}
137
+ return cls.from_lru_numpy(
138
+ nu_log=p['nu_log'], theta_log=p['theta_log'],
139
+ B_re=p['B_re'], B_im=p['B_im'], gamma_log=p['gamma_log'],
140
+ C_re=p['C_re'], C_im=p['C_im'], D=p['D'],
141
+ )
142
+
143
+
144
+ @dataclass(frozen=True)
145
+ class ReducedSSM:
146
+ A: np.ndarray
147
+ B: np.ndarray
148
+ C: np.ndarray
149
+ D: np.ndarray
150
+ r: int
151
+ error_bound: float
152
+
153
+
154
+ @dataclass(frozen=True)
155
+ class ReducedDiagonalSSM:
156
+ lam: np.ndarray
157
+ B: np.ndarray
158
+ C: np.ndarray
159
+ D: np.ndarray
160
+ r: int
161
+ error_bound: float
162
+
163
+
164
+ @dataclass(frozen=True)
165
+ class Balancing:
166
+ Lp: np.ndarray
167
+ Lq: np.ndarray
168
+ U: np.ndarray
169
+ sigma: np.ndarray
170
+ Vh: np.ndarray
171
+
172
+
173
+ def tensor_to_numpy(x: torch.Tensor) -> np.ndarray:
174
+ return x.detach().cpu().numpy().copy()
175
+
176
+
177
+ class CompreSSM:
178
+ @staticmethod
179
+ def calculate_gramians_diagonal(lam: np.ndarray, *, B: np.ndarray, C: np.ndarray):
180
+ denom = 1.0 - lam[:, None] * lam.conj()[None, :]
181
+ P = (B @ B.conj().T) / denom
182
+ Q = (C.conj().T @ C) / denom.conj()
183
+ return (P + P.conj().T) / 2, (Q + Q.conj().T) / 2
184
+
185
+ @staticmethod
186
+ def calculate_gramians(A: np.ndarray, *, B: np.ndarray, C: np.ndarray):
187
+ P = solve_discrete_lyapunov(A, B @ B.conj().T)
188
+ Q = solve_discrete_lyapunov(A.conj().T, C.conj().T @ C)
189
+ return (P + P.conj().T) / 2, (Q + Q.conj().T) / 2
190
+
191
+ @staticmethod
192
+ def diagonalize_reduced(Ar: np.ndarray, Br: np.ndarray, Cr: np.ndarray):
193
+ lam_r, V = np.linalg.eig(Ar)
194
+ return lam_r, np.linalg.solve(V, Br), Cr @ V
195
+
196
+ @staticmethod
197
+ def cholesky_with_jitter(M: np.ndarray, *, name: str, jitter: float = 1e-12, max_tries: int = 8):
198
+ n = M.shape[0]
199
+ scale = float(np.real(np.trace(M))) / n
200
+ for k in range(max_tries):
201
+ eps = 0.0 if k == 0 else jitter * scale * 10.0 ** (k - 1)
202
+ try:
203
+ return cholesky(M + eps * np.eye(n), lower=True)
204
+ except np.linalg.LinAlgError:
205
+ continue
206
+ raise np.linalg.LinAlgError(f"Gramian not positive definite on matrix {name}, even with jitter")
207
+
208
+ @staticmethod
209
+ def calculate_sqrt(P: np.ndarray, *, Q: np.ndarray):
210
+ Lp = CompreSSM.cholesky_with_jitter(P, name="P (controllability Gramian)")
211
+ Lq = CompreSSM.cholesky_with_jitter(Q, name="Q (observability Gramian)")
212
+ return Lp, Lq
213
+
214
+ @staticmethod
215
+ def HSVs(Lp: np.ndarray, *, Lq: np.ndarray):
216
+ return svd(Lq.conj().T @ Lp)
217
+
218
+ @staticmethod
219
+ def balance(P: np.ndarray, *, Q: np.ndarray) -> Balancing:
220
+ Lp, Lq = CompreSSM.calculate_sqrt(P, Q=Q)
221
+ U, sigma, Vh = CompreSSM.HSVs(Lp, Lq=Lq)
222
+ return Balancing(Lp=Lp, Lq=Lq, U=U, sigma=sigma, Vh=Vh)
223
+
224
+ @staticmethod
225
+ def rank_selection(sigma: np.ndarray, *, threshold: float) -> int:
226
+ """
227
+ Finds the smallest rank retaining the requested fraction of the HSV sum.
228
+
229
+ Args:
230
+ sigma: 1D array of Hankel singular values sorted in descending order.
231
+ threshold: Fraction of the total HSV sum to retain (1.0 - tol).
232
+ """
233
+ sigma = np.asarray(sigma)
234
+ if sigma.ndim != 1 or not np.all(np.isfinite(sigma)) or np.any(sigma < 0):
235
+ raise ValueError("sigma must be a one-dimensional array of finite non-negative values")
236
+ if not np.isfinite(threshold) or not 0 <= threshold <= 1:
237
+ raise ValueError("threshold must be between zero and one")
238
+ if sigma.size == 0:
239
+ return 0
240
+
241
+ total = np.sum(sigma)
242
+ if total <= 0:
243
+ return len(sigma)
244
+ cumulative_sum = np.cumsum(sigma)
245
+ ranks = np.flatnonzero(cumulative_sum >= threshold * total)
246
+ return int(ranks[0] + 1) if ranks.size else len(sigma)
247
+
248
+ @staticmethod
249
+ def should_compress(r: int, n: int, *, max_ratio: float):
250
+ return 0 < r < max_ratio * n
251
+
252
+ @staticmethod
253
+ def calculate_transform_matrix(Lp: np.ndarray, *, r: int, sigma: np.ndarray, Vh: np.ndarray):
254
+ if np.any(sigma[:r] <= 0):
255
+ raise np.linalg.LinAlgError("Retained Hankel singular values must be positive")
256
+ Vr = Vh[:r].conj().T
257
+ diag_entries = sigma[:r] ** -0.5
258
+ return Lp @ Vr @ np.diag(diag_entries), diag_entries
259
+
260
+ @staticmethod
261
+ def calculate_projection_matrix(Lq: np.ndarray, *, r: int, diag_entries: np.ndarray, U: np.ndarray):
262
+ return np.diag(diag_entries) @ U[:, :r].conj().T @ Lq.conj().T
263
+
264
+ @staticmethod
265
+ def truncate_A(A: np.ndarray, *, S: np.ndarray, T: np.ndarray):
266
+ return S @ A @ T
267
+
268
+ @staticmethod
269
+ def truncate_B(S: np.ndarray, *, B: np.ndarray):
270
+ return S @ B
271
+
272
+ @staticmethod
273
+ def truncate_C(C: np.ndarray, *, T: np.ndarray):
274
+ return C @ T
275
+
276
+ @staticmethod
277
+ def truncate(A: np.ndarray, B: np.ndarray, C: np.ndarray, *, balancing: Balancing, r: int):
278
+ T, diag_entries = CompreSSM.calculate_transform_matrix(
279
+ balancing.Lp, r=r, sigma=balancing.sigma, Vh=balancing.Vh
280
+ )
281
+ S = CompreSSM.calculate_projection_matrix(
282
+ balancing.Lq, r=r, diag_entries=diag_entries, U=balancing.U
283
+ )
284
+ return (
285
+ CompreSSM.truncate_A(A, S=S, T=T),
286
+ CompreSSM.truncate_B(S, B=B),
287
+ CompreSSM.truncate_C(C, T=T),
288
+ )
289
+
290
+
291
+ def compress_from_gramians(A: np.ndarray, B: np.ndarray, C: np.ndarray, *, P: np.ndarray, Q: np.ndarray, config: CompressionConfig):
292
+ balancing = CompreSSM.balance(P, Q=Q)
293
+ r = CompreSSM.rank_selection(balancing.sigma, threshold=config.threshold)
294
+
295
+ if not CompreSSM.should_compress(r, A.shape[0], max_ratio=config.max_ratio):
296
+ return None
297
+
298
+ Ar, Br, Cr = CompreSSM.truncate(A, B, C, balancing=balancing, r=r)
299
+ error_bound = 2 * float(np.real(np.sum(balancing.sigma[r:])))
300
+ return Ar, Br, Cr, r, error_bound
301
+
302
+
303
+ def compress_numpy(system: DiscreteSSM, *, config: CompressionConfig) -> ReducedSSM | None:
304
+ """Apply discrete-time balanced truncation to one stable, minimal LTI system.
305
+
306
+ The returned error bound is the standard H-infinity bound for balanced
307
+ truncation; it assumes the LTI system is stable and minimal.
308
+ """
309
+ P, Q = CompreSSM.calculate_gramians(system.A, B=system.B, C=system.C)
310
+ result = compress_from_gramians(system.A, system.B, system.C, P=P, Q=Q, config=config)
311
+ if result is None:
312
+ return None
313
+ Ar, Br, Cr, r, error_bound = result
314
+ return ReducedSSM(A=Ar, B=Br, C=Cr, D=system.D, r=r, error_bound=error_bound)
315
+
316
+
317
+ def compress_diagonal_numpy(system: DiagonalDiscreteSSM, *, config: CompressionConfig) -> ReducedDiagonalSSM | None:
318
+ """Balanced truncation of a diagonal system, returned in diagonal form."""
319
+ P, Q = CompreSSM.calculate_gramians_diagonal(system.lam, B=system.B, C=system.C)
320
+ result = compress_from_gramians(np.diag(system.lam), system.B, system.C, P=P, Q=Q, config=config)
321
+ if result is None:
322
+ return None
323
+ Ar, Br, Cr, r, error_bound = result
324
+ lam_r, Br_diag, Cr_diag = CompreSSM.diagonalize_reduced(Ar, Br, Cr)
325
+ return ReducedDiagonalSSM(lam=lam_r, B=Br_diag, C=Cr_diag, D=system.D, r=r, error_bound=error_bound)
326
+
327
+
328
+ def compress_module(module: nn.Module, *, config: CompressionConfig) -> ReducedSSM | None:
329
+ """Extract and compress the expected fixed-step diagonal SSM; does not mutate the module."""
330
+ system = DiscreteSSM.from_module(module, dt=config.dt)
331
+ return compress_numpy(system, config=config)
332
+
333
+
334
+ def compress_lru_module(module: nn.Module, *, config: CompressionConfig) -> ReducedDiagonalSSM | None:
335
+ """Extract and compress one LRU block; does not mutate the module."""
336
+ system = DiagonalDiscreteSSM.from_lru_module(module)
337
+ return compress_diagonal_numpy(system, config=config)
338
+
339
+
340
+ def lru_arrays_from_diagonal(reduced: ReducedDiagonalSSM) -> dict[str, np.ndarray]:
341
+ modulus = np.clip(np.abs(reduced.lam), 1e-12, 1.0 - 1e-12)
342
+ theta = np.maximum(np.angle(reduced.lam) % (2 * np.pi), 1e-8)
343
+ gamma = np.sqrt(1.0 - modulus ** 2)
344
+ B_stored = reduced.B / gamma[:, None]
345
+ return {
346
+ 'nu_log': np.log(-np.log(modulus)),
347
+ 'theta_log': np.log(theta),
348
+ 'gamma_log': np.log(gamma),
349
+ 'B_re': B_stored.real,
350
+ 'B_im': B_stored.imag,
351
+ 'C_re': reduced.C.real,
352
+ 'C_im': reduced.C.imag,
353
+ }
354
+
355
+
356
+ def replace_parameter(module: nn.Module, name: str, array: np.ndarray) -> None:
357
+ old = getattr(module, name)
358
+ new = torch.tensor(np.ascontiguousarray(array), dtype=old.dtype, device=old.device)
359
+ setattr(module, name, nn.Parameter(new, requires_grad=old.requires_grad))
360
+
361
+
362
+ def write_back_lru(module: nn.Module, *, reduced: ReducedDiagonalSSM) -> None:
363
+ """Replace the LRU block's parameters with the reduced ones. D is unchanged."""
364
+ for name, array in lru_arrays_from_diagonal(reduced).items():
365
+ replace_parameter(module, name, array)
366
+
367
+
368
+ def find_blocks(model: nn.Module, *, block_cls: type) -> list[nn.Module]:
369
+ return [m for m in model.modules() if isinstance(m, block_cls)]
370
+
371
+
372
+ def reduce_lru_blocks(blocks: Iterable[nn.Module], *, config: CompressionConfig) -> list[ReducedDiagonalSSM | None]:
373
+ results = []
374
+ for block in blocks:
375
+ reduced = compress_lru_module(block, config=config)
376
+ if reduced is not None:
377
+ write_back_lru(block, reduced=reduced)
378
+ results.append(reduced)
379
+ return results
380
+
381
+
382
+ def run_scheduled_reduction(
383
+ *,
384
+ step: int,
385
+ schedule: ReductionSchedule,
386
+ blocks: Iterable[nn.Module],
387
+ config: CompressionConfig,
388
+ make_optimizer: Callable[[], torch.optim.Optimizer],
389
+ ):
390
+ """Returns a freshly built optimizer if any block was reduced, otherwise None."""
391
+ if not schedule.fires_at(step):
392
+ return None
393
+ results = reduce_lru_blocks(blocks, config=config)
394
+ if all(result is None for result in results):
395
+ return None
396
+ return make_optimizer()