Skip to content

Factored Matrix

FactoredMatrix represents a matrix as A @ B. This mirrors the common TransformerLens pattern for composing projection matrices without immediately materializing every dense product.

Supported operations:

  • Dense product through .AB and reverse product through .BA.
  • Transpose through .T.
  • Matrix/vector multiplication with @; left matrix multiplication preserves a factored result, right matrix multiplication preserves a factored result, and vector multiplication returns a vector.
  • Scalar multiplication with *.
  • Factored composition with another FactoredMatrix.
  • Frobenius norm with .norm(); leading dimensions are preserved.
  • SVD-backed .U, .S, .V, .Vh, .collapse_l(), .collapse_r(), .make_even(), and .eigenvalues when NumPy compatible data is available.
  • .svd() returns (U, S, V), with .Vh retained as a TransformerLens compatibility alias for .V.
  • Torch tensor inputs keep tensor outputs for SVD, collapse, even-factor, norm, and composition-score workflows.
  • Rectangular matrix eigenvalues are computed from BA, matching TransformerLens' non-zero spectrum convention.
  • unsqueeze(dim) and squeeze(dim) for leading-dimension shape management.
  • TransformerLens-style indexing over leading, row, and column dimensions while preserving a factored matrix result.
  • Singleton leading dimensions are broadcast on construction, matching TransformerLens' layer/head batching behavior.
  • composition_scores(left, right) for QK/OV-style circuit composition analysis, with optional pairwise leading-dimension broadcasting.
  • Lightweight list and tensor-like fallbacks for projects that do not install TransformerLens.

Example:

from SafeLens.core.factored_matrix import FactoredMatrix, composition_scores

W_QK = FactoredMatrix([[1, 2], [3, 4]], [[2, 0], [0, 2]])

assert W_QK.AB == [[2.0, 4.0], [6.0, 8.0]]
assert W_QK @ [1, 1] == [6.0, 14.0]
assert 0 <= composition_scores(W_QK, W_QK) <= 1

Low-rank factored matrix utilities inspired by TransformerLens.

FactoredMatrix dataclass

Represent a matrix as a product A @ B.

Source code in src/SafeLens/core/factored_matrix.py 
 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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
@dataclass(frozen=True)
class FactoredMatrix:
    """Represent a matrix as a product `A @ B`."""

    A: Any
    B: Any

    def __post_init__(self) -> None:
        left_shape = shape_of(self.A)
        right_shape = shape_of(self.B)
        if len(left_shape) < 2 or len(right_shape) < 2:
            raise ValueError(
                "FactoredMatrix factors must be at least rank-2, "
                f"got {left_shape} and {right_shape}."
            )
        if left_shape[-1] != right_shape[-2]:
            raise ValueError(
                f"FactoredMatrix inner dimensions must match, got {left_shape} and {right_shape}."
            )
        leading_shape = broadcast_leading_shape(left_shape[:-2], right_shape[:-2])
        object.__setattr__(self, "A", broadcast_to_shape(self.A, leading_shape + left_shape[-2:]))
        object.__setattr__(self, "B", broadcast_to_shape(self.B, leading_shape + right_shape[-2:]))

    @property
    def pair(self) -> tuple[Any, Any]:
        """Return the matrix factors."""
        return self.A, self.B

    @property
    def AB(self) -> Any:
        """Return the dense product `A @ B`."""
        return matmul(self.A, self.B)

    @property
    def BA(self) -> Any:
        """Return the reverse dense product `B @ A`."""
        if self.ldim != self.rdim:
            raise ValueError(f"Can only take BA if ldim == rdim, got dense shape {self.shape}.")
        return matmul(self.B, self.A)

    @property
    def T(self) -> FactoredMatrix:
        """Return the transposed factored matrix."""
        return FactoredMatrix(transpose(self.B), transpose(self.A))

    @property
    def ndim(self) -> int:
        """Return the rank of the dense product."""
        return len(self.shape)

    @property
    def ldim(self) -> int:
        """Return the row dimension of the dense product."""
        return self.shape[-2]

    @property
    def rdim(self) -> int:
        """Return the column dimension of the dense product."""
        return self.shape[-1]

    @property
    def mdim(self) -> int:
        """Return the hidden dimension shared by the two factors."""
        return shape_of(self.B)[-2]

    @property
    def has_leading_dims(self) -> bool:
        """Return whether either factor has leading batch-like dimensions."""
        return len(shape_of(self.A)) > 2 or len(shape_of(self.B)) > 2

    @property
    def shape(self) -> tuple[int, ...]:
        """Return the dense matrix shape."""
        left_shape = shape_of(self.A)
        right_shape = shape_of(self.B)
        leading_shape = broadcast_leading_shape(left_shape[:-2], right_shape[:-2])
        return leading_shape + (left_shape[-2], right_shape[-1])

    @property
    def U(self) -> Any:
        """Return left singular vectors from `svd()`."""
        return self.svd()[0]

    @property
    def S(self) -> Any:
        """Return singular values from `svd()`."""
        return self.svd()[1]

    @property
    def Vh(self) -> Any:
        """Deprecated alias for right singular vectors from `svd()`."""
        warnings.warn(
            "FactoredMatrix.Vh returns V (right singular vectors), not Vh. "
            "Use .V for the canonical name.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.svd()[2]

    @property
    def V(self) -> Any:
        """Return right singular vectors from `svd()`."""
        return self.svd()[2]

    @property
    def eigenvalues(self) -> Any:
        """Return eigenvalues of `BA`, matching the non-zero spectrum of `AB`."""
        input_matrix = self.BA
        try:
            import torch

            if isinstance(input_matrix, torch.Tensor):
                if input_matrix.dtype in {torch.bfloat16, torch.float16}:
                    input_matrix = input_matrix.to(torch.float32)
                return torch.linalg.eig(input_matrix).eigenvalues
        except ImportError:
            pass
        try:
            import numpy as np

            eigenvalues = np.linalg.eigvals(np.asarray(input_matrix))
            return _real_eigenvalues_if_close(eigenvalues.tolist())
        except Exception:
            try:
                return eigenvalues_nested(input_matrix)
            except Exception as fallback_exc:
                raise RuntimeError(
                    "FactoredMatrix.eigenvalues requires numpy-compatible data."
                ) from fallback_exc

    def __matmul__(self, other: Any) -> Any:
        """Multiply by another matrix/vector/factored matrix."""
        if isinstance(other, FactoredMatrix):
            return (self @ other.A) @ other.B
        if _is_vector_like(other):
            return matmul(self.AB, other)
        right_shape = shape_of(other)
        if len(right_shape) >= 2 and right_shape[-2] == self.shape[-1]:
            if self.rdim > self.mdim:
                return FactoredMatrix(self.A, matmul(self.B, other))
            return FactoredMatrix(self.AB, other)
        return matmul(self.AB, other)

    def __rmatmul__(self, other: Any) -> Any:
        """Left multiply the dense product."""
        if isinstance(other, FactoredMatrix):
            return other @ self
        if _is_vector_like(other):
            return matmul(other, self.AB)
        left_shape = shape_of(other)
        if len(left_shape) >= 2 and left_shape[-1] == self.shape[-2]:
            if self.ldim > self.mdim:
                return FactoredMatrix(matmul(other, self.A), self.B)
            return FactoredMatrix(other, self.AB)
        return matmul(other, self.AB)

    def __mul__(self, scalar: Any) -> FactoredMatrix:
        """Scale the factored matrix by multiplying the left factor."""
        if not _is_scalar_like(scalar):
            raise TypeError("FactoredMatrix scalar multiplication expects a scalar.")
        return FactoredMatrix(scale_values(self.A, scalar), self.B)

    def __rmul__(self, scalar: Any) -> FactoredMatrix:
        """Right scalar multiplication."""
        return self * scalar

    def __getitem__(self, index: Any) -> FactoredMatrix:
        """Index leading, row, and column dimensions while preserving factored rank."""
        normalized = expand_ellipsis(normalize_index(index), len(self.shape))
        indexed_dims = len([item for item in normalized if item is not None])
        leading_dims = max(0, len(self.shape) - 2)
        if indexed_dims <= leading_dims:
            return FactoredMatrix(index_value(self.A, normalized), index_value(self.B, normalized))
        if indexed_dims == leading_dims + 1:
            row_index = convert_int_to_slice(normalized, -1)
            return FactoredMatrix(
                index_value(self.A, row_index), index_value(self.B, row_index[:-1])
            )
        if indexed_dims == leading_dims + 2:
            row_col_index = convert_int_to_slice(convert_int_to_slice(normalized, -1), -2)
            return FactoredMatrix(
                index_value(self.A, row_col_index[:-1]),
                index_value(self.B, row_col_index[:-2] + (FULL_SLICE, row_col_index[-1])),
            )
        raise ValueError(
            f"{normalized!r} is too long an index for a FactoredMatrix with shape {self.shape}."
        )

    def __repr__(self) -> str:
        """Return a TransformerLens-style summary."""
        return f"FactoredMatrix: Shape({self.shape}), Hidden Dim({shape_of(self.B)[-2]})"

    def to_dense_right(self, right: Any) -> Any:
        """Return dense `(A @ B) @ right`."""
        return matmul(self.AB, right)

    def svd(self) -> tuple[Any, Any, Any]:
        """Return `(U, S, V)` with right singular vectors in `V`, matching TransformerLens."""
        try:
            import torch

            if isinstance(self.A, torch.Tensor) or isinstance(self.B, torch.Tensor):
                a = (
                    self.A
                    if isinstance(self.A, torch.Tensor)
                    else torch.as_tensor(
                        self.A,
                        dtype=getattr(self.B, "dtype", None),
                        device=getattr(self.B, "device", None),
                    )
                )
                b = (
                    self.B
                    if isinstance(self.B, torch.Tensor)
                    else torch.as_tensor(
                        self.B,
                        dtype=a.dtype,
                        device=a.device,
                    )
                )
                if not torch.is_floating_point(a) or a.dtype in {torch.bfloat16, torch.float16}:
                    a = a.to(torch.float32)
                if not torch.is_floating_point(b) or b.dtype in {torch.bfloat16, torch.float16}:
                    b = b.to(torch.float32)
                u_a, s_a, vh_a = torch.linalg.svd(a, full_matrices=False)
                u_b, s_b, vh_b = torch.linalg.svd(b, full_matrices=False)
                v_a = vh_a.transpose(-1, -2)
                v_b = vh_b.transpose(-1, -2)
                middle = s_a[..., :, None] * (v_a.transpose(-1, -2) @ u_b) * s_b[..., None, :]
                u_m, s_m, vh_m = torch.linalg.svd(middle, full_matrices=False)
                v_m = vh_m.transpose(-1, -2)
                return u_a @ u_m, s_m, v_b @ v_m
        except Exception:
            pass
        try:
            import numpy as np

            a = _promote_numpy_svd_array(np.asarray(self.A), np)
            b = _promote_numpy_svd_array(np.asarray(self.B), np)
            u_a, s_a, vh_a = np.linalg.svd(a, full_matrices=False)
            u_b, s_b, vh_b = np.linalg.svd(b, full_matrices=False)
            v_a = np.swapaxes(vh_a, -1, -2)
            v_b = np.swapaxes(vh_b, -1, -2)
            middle = s_a[..., :, None] * (np.swapaxes(v_a, -1, -2) @ u_b) * s_b[..., None, :]
            u_m, s_m, vh_m = np.linalg.svd(middle, full_matrices=False)
            v_m = np.swapaxes(vh_m, -1, -2)
            return (u_a @ u_m).tolist(), s_m.tolist(), (v_b @ v_m).tolist()
        except Exception:
            try:
                return svd_nested(self.AB)
            except Exception as fallback_exc:
                raise RuntimeError(
                    "FactoredMatrix.svd requires numpy-compatible data."
                ) from fallback_exc

    def norm(self) -> float:
        """Return the Frobenius norm without materializing the dense product."""
        return factored_frobenius_norm(self.A, self.B)

    def collapse_l(self) -> Any:
        """Collapse the left orthogonal factor using the SVD."""
        return scale_rows(transpose(self.V), self.S)

    def collapse_r(self) -> Any:
        """Collapse the right orthogonal factor using the SVD."""
        return scale_columns(self.U, self.S)

    def make_even(self) -> FactoredMatrix:
        """Return an equivalent more balanced factorization using SVD."""
        left, right = _make_even_factors(*self.svd())
        return FactoredMatrix(left, right)

    def get_corner(self, k: int = 3) -> Any:
        """Return the top-left dense corner."""
        left_index = expand_ellipsis((Ellipsis, slice(None, k), FULL_SLICE), len(shape_of(self.A)))
        right_index = expand_ellipsis((Ellipsis, FULL_SLICE, slice(None, k)), len(shape_of(self.B)))
        return matrix_corner(
            matmul(index_value(self.A, left_index), index_value(self.B, right_index)), k
        )

    def unsqueeze(self, dim: int | None = None, *, k: int | None = None) -> FactoredMatrix:
        """Add a leading dimension to both factors."""
        if dim is None:
            if k is None:
                raise TypeError("unsqueeze() missing required argument: 'dim' or 'k'")
            dim = k
        elif k is not None:
            raise TypeError("Pass only one of `dim` or `k`.")
        return FactoredMatrix(unsqueeze_dim(self.A, dim), unsqueeze_dim(self.B, dim))

    def squeeze(self, dim: int | None = None) -> FactoredMatrix:
        """Remove singleton dimensions from both factors."""
        return FactoredMatrix(squeeze_dim(self.A, dim), squeeze_dim(self.B, dim))

AB property

Return the dense product A @ B.

BA property

Return the reverse dense product B @ A.

S property

Return singular values from svd().

T property

Return the transposed factored matrix.

U property

Return left singular vectors from svd().

V property

Return right singular vectors from svd().

Vh property

Deprecated alias for right singular vectors from svd().

eigenvalues property

Return eigenvalues of BA, matching the non-zero spectrum of AB.

has_leading_dims property

Return whether either factor has leading batch-like dimensions.

ldim property

Return the row dimension of the dense product.

mdim property

Return the hidden dimension shared by the two factors.

ndim property

Return the rank of the dense product.

pair property

Return the matrix factors.

rdim property

Return the column dimension of the dense product.

shape property

Return the dense matrix shape.

__getitem__(index)

Index leading, row, and column dimensions while preserving factored rank.

Source code in src/SafeLens/core/factored_matrix.py 
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def __getitem__(self, index: Any) -> FactoredMatrix:
    """Index leading, row, and column dimensions while preserving factored rank."""
    normalized = expand_ellipsis(normalize_index(index), len(self.shape))
    indexed_dims = len([item for item in normalized if item is not None])
    leading_dims = max(0, len(self.shape) - 2)
    if indexed_dims <= leading_dims:
        return FactoredMatrix(index_value(self.A, normalized), index_value(self.B, normalized))
    if indexed_dims == leading_dims + 1:
        row_index = convert_int_to_slice(normalized, -1)
        return FactoredMatrix(
            index_value(self.A, row_index), index_value(self.B, row_index[:-1])
        )
    if indexed_dims == leading_dims + 2:
        row_col_index = convert_int_to_slice(convert_int_to_slice(normalized, -1), -2)
        return FactoredMatrix(
            index_value(self.A, row_col_index[:-1]),
            index_value(self.B, row_col_index[:-2] + (FULL_SLICE, row_col_index[-1])),
        )
    raise ValueError(
        f"{normalized!r} is too long an index for a FactoredMatrix with shape {self.shape}."
    )

__matmul__(other)

Multiply by another matrix/vector/factored matrix.

Source code in src/SafeLens/core/factored_matrix.py 
149
150
151
152
153
154
155
156
157
158
159
160
def __matmul__(self, other: Any) -> Any:
    """Multiply by another matrix/vector/factored matrix."""
    if isinstance(other, FactoredMatrix):
        return (self @ other.A) @ other.B
    if _is_vector_like(other):
        return matmul(self.AB, other)
    right_shape = shape_of(other)
    if len(right_shape) >= 2 and right_shape[-2] == self.shape[-1]:
        if self.rdim > self.mdim:
            return FactoredMatrix(self.A, matmul(self.B, other))
        return FactoredMatrix(self.AB, other)
    return matmul(self.AB, other)

__mul__(scalar)

Scale the factored matrix by multiplying the left factor.

Source code in src/SafeLens/core/factored_matrix.py 
175
176
177
178
179
def __mul__(self, scalar: Any) -> FactoredMatrix:
    """Scale the factored matrix by multiplying the left factor."""
    if not _is_scalar_like(scalar):
        raise TypeError("FactoredMatrix scalar multiplication expects a scalar.")
    return FactoredMatrix(scale_values(self.A, scalar), self.B)

__repr__()

Return a TransformerLens-style summary.

Source code in src/SafeLens/core/factored_matrix.py 
207
208
209
def __repr__(self) -> str:
    """Return a TransformerLens-style summary."""
    return f"FactoredMatrix: Shape({self.shape}), Hidden Dim({shape_of(self.B)[-2]})"

__rmatmul__(other)

Left multiply the dense product.

Source code in src/SafeLens/core/factored_matrix.py 
162
163
164
165
166
167
168
169
170
171
172
173
def __rmatmul__(self, other: Any) -> Any:
    """Left multiply the dense product."""
    if isinstance(other, FactoredMatrix):
        return other @ self
    if _is_vector_like(other):
        return matmul(other, self.AB)
    left_shape = shape_of(other)
    if len(left_shape) >= 2 and left_shape[-1] == self.shape[-2]:
        if self.ldim > self.mdim:
            return FactoredMatrix(matmul(other, self.A), self.B)
        return FactoredMatrix(other, self.AB)
    return matmul(other, self.AB)

__rmul__(scalar)

Right scalar multiplication.

Source code in src/SafeLens/core/factored_matrix.py 
181
182
183
def __rmul__(self, scalar: Any) -> FactoredMatrix:
    """Right scalar multiplication."""
    return self * scalar

collapse_l()

Collapse the left orthogonal factor using the SVD.

Source code in src/SafeLens/core/factored_matrix.py 
278
279
280
def collapse_l(self) -> Any:
    """Collapse the left orthogonal factor using the SVD."""
    return scale_rows(transpose(self.V), self.S)

collapse_r()

Collapse the right orthogonal factor using the SVD.

Source code in src/SafeLens/core/factored_matrix.py 
282
283
284
def collapse_r(self) -> Any:
    """Collapse the right orthogonal factor using the SVD."""
    return scale_columns(self.U, self.S)

get_corner(k=3)

Return the top-left dense corner.

Source code in src/SafeLens/core/factored_matrix.py 
291
292
293
294
295
296
297
def get_corner(self, k: int = 3) -> Any:
    """Return the top-left dense corner."""
    left_index = expand_ellipsis((Ellipsis, slice(None, k), FULL_SLICE), len(shape_of(self.A)))
    right_index = expand_ellipsis((Ellipsis, FULL_SLICE, slice(None, k)), len(shape_of(self.B)))
    return matrix_corner(
        matmul(index_value(self.A, left_index), index_value(self.B, right_index)), k
    )

make_even()

Return an equivalent more balanced factorization using SVD.

Source code in src/SafeLens/core/factored_matrix.py 
286
287
288
289
def make_even(self) -> FactoredMatrix:
    """Return an equivalent more balanced factorization using SVD."""
    left, right = _make_even_factors(*self.svd())
    return FactoredMatrix(left, right)

norm()

Return the Frobenius norm without materializing the dense product.

Source code in src/SafeLens/core/factored_matrix.py 
274
275
276
def norm(self) -> float:
    """Return the Frobenius norm without materializing the dense product."""
    return factored_frobenius_norm(self.A, self.B)

squeeze(dim=None)

Remove singleton dimensions from both factors.

Source code in src/SafeLens/core/factored_matrix.py 
309
310
311
def squeeze(self, dim: int | None = None) -> FactoredMatrix:
    """Remove singleton dimensions from both factors."""
    return FactoredMatrix(squeeze_dim(self.A, dim), squeeze_dim(self.B, dim))

svd()

Return (U, S, V) with right singular vectors in V, matching TransformerLens.

Source code in src/SafeLens/core/factored_matrix.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
def svd(self) -> tuple[Any, Any, Any]:
    """Return `(U, S, V)` with right singular vectors in `V`, matching TransformerLens."""
    try:
        import torch

        if isinstance(self.A, torch.Tensor) or isinstance(self.B, torch.Tensor):
            a = (
                self.A
                if isinstance(self.A, torch.Tensor)
                else torch.as_tensor(
                    self.A,
                    dtype=getattr(self.B, "dtype", None),
                    device=getattr(self.B, "device", None),
                )
            )
            b = (
                self.B
                if isinstance(self.B, torch.Tensor)
                else torch.as_tensor(
                    self.B,
                    dtype=a.dtype,
                    device=a.device,
                )
            )
            if not torch.is_floating_point(a) or a.dtype in {torch.bfloat16, torch.float16}:
                a = a.to(torch.float32)
            if not torch.is_floating_point(b) or b.dtype in {torch.bfloat16, torch.float16}:
                b = b.to(torch.float32)
            u_a, s_a, vh_a = torch.linalg.svd(a, full_matrices=False)
            u_b, s_b, vh_b = torch.linalg.svd(b, full_matrices=False)
            v_a = vh_a.transpose(-1, -2)
            v_b = vh_b.transpose(-1, -2)
            middle = s_a[..., :, None] * (v_a.transpose(-1, -2) @ u_b) * s_b[..., None, :]
            u_m, s_m, vh_m = torch.linalg.svd(middle, full_matrices=False)
            v_m = vh_m.transpose(-1, -2)
            return u_a @ u_m, s_m, v_b @ v_m
    except Exception:
        pass
    try:
        import numpy as np

        a = _promote_numpy_svd_array(np.asarray(self.A), np)
        b = _promote_numpy_svd_array(np.asarray(self.B), np)
        u_a, s_a, vh_a = np.linalg.svd(a, full_matrices=False)
        u_b, s_b, vh_b = np.linalg.svd(b, full_matrices=False)
        v_a = np.swapaxes(vh_a, -1, -2)
        v_b = np.swapaxes(vh_b, -1, -2)
        middle = s_a[..., :, None] * (np.swapaxes(v_a, -1, -2) @ u_b) * s_b[..., None, :]
        u_m, s_m, vh_m = np.linalg.svd(middle, full_matrices=False)
        v_m = np.swapaxes(vh_m, -1, -2)
        return (u_a @ u_m).tolist(), s_m.tolist(), (v_b @ v_m).tolist()
    except Exception:
        try:
            return svd_nested(self.AB)
        except Exception as fallback_exc:
            raise RuntimeError(
                "FactoredMatrix.svd requires numpy-compatible data."
            ) from fallback_exc

to_dense_right(right)

Return dense (A @ B) @ right.

Source code in src/SafeLens/core/factored_matrix.py 
211
212
213
def to_dense_right(self, right: Any) -> Any:
    """Return dense `(A @ B) @ right`."""
    return matmul(self.AB, right)

unsqueeze(dim=None, *, k=None)

Add a leading dimension to both factors.

Source code in src/SafeLens/core/factored_matrix.py 
299
300
301
302
303
304
305
306
307
def unsqueeze(self, dim: int | None = None, *, k: int | None = None) -> FactoredMatrix:
    """Add a leading dimension to both factors."""
    if dim is None:
        if k is None:
            raise TypeError("unsqueeze() missing required argument: 'dim' or 'k'")
        dim = k
    elif k is not None:
        raise TypeError("Pass only one of `dim` or `k`.")
    return FactoredMatrix(unsqueeze_dim(self.A, dim), unsqueeze_dim(self.B, dim))

as_2d(value)

Convert a tensor-like value to a 2D Python list.

Source code in src/SafeLens/core/factored_matrix.py 
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
def as_2d(value: Any) -> list[list[Any]]:
    """Convert a tensor-like value to a 2D Python list."""
    tolist = getattr(value, "tolist", None)
    if callable(tolist):
        value = tolist()
    if not _is_sequence(value):
        return [[value]]
    if not value:
        return []
    if _is_sequence(value[0]):
        return [list(row) for row in value]
    return [list(value)]

broadcast_pair(left, right)

Broadcast two tensor-like or nested values to a shared shape.

Source code in src/SafeLens/core/factored_matrix.py 
566
567
568
569
def broadcast_pair(left: Any, right: Any) -> tuple[Any, Any]:
    """Broadcast two tensor-like or nested values to a shared shape."""
    target_shape = broadcast_leading_shape(shape_of(left), shape_of(right))
    return broadcast_to_shape(left, target_shape), broadcast_to_shape(right, target_shape)

composition_scores(left, right, broadcast_dims=True)

Return TransformerLens-style composition scores for two factored matrices.

Source code in src/SafeLens/core/factored_matrix.py 
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
def composition_scores(
    left: FactoredMatrix,
    right: FactoredMatrix,
    broadcast_dims: bool = True,
) -> Any:
    """Return TransformerLens-style composition scores for two factored matrices."""
    if broadcast_dims:
        left_leading = left.ndim - 2
        right_leading = right.ndim - 2
        for dim in range(left_leading):
            right = right.unsqueeze(dim)
        for dim in range(right_leading):
            left = left.unsqueeze(dim + left_leading)
    if left.rdim != right.ldim:
        raise ValueError(
            "Composition scores require left.rdim == right.ldim, "
            f"got left shape {left.shape} and right shape {right.shape}."
        )

    collapsed_left = left.collapse_l()
    collapsed_right = right.collapse_r()
    composed = matmul(collapsed_left, collapsed_right)
    denominator = multiply_values(frobenius_norm(collapsed_left), frobenius_norm(collapsed_right))
    return divide_values(frobenius_norm(composed), denominator)

divide_values(left, right)

Divide tensor-like or nested values elementwise.

Source code in src/SafeLens/core/factored_matrix.py 
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
def divide_values(left: Any, right: Any) -> Any:
    """Divide tensor-like or nested values elementwise."""
    try:
        return left / right
    except Exception:
        pass
    backend_result = _try_backend_elementwise(left, right, op="divide")
    if backend_result is not _BACKEND_MATMUL_NOT_AVAILABLE:
        return backend_result
    left_shape = shape_of(left)
    right_shape = shape_of(right)
    if left_shape or right_shape:
        left_b, right_b = broadcast_pair(left, right)
        if _is_sequence(left_b) and _is_sequence(right_b):
            return [
                divide_values(left_item, right_item)
                for left_item, right_item in zip(left_b, right_b, strict=True)
            ]
    return float(left) / float(right)

eigenvalues_nested(matrix)

Pure-Python eigenvalue fallback for small real matrices.

Source code in src/SafeLens/core/factored_matrix.py 
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
def eigenvalues_nested(matrix: Any) -> Any:
    """Pure-Python eigenvalue fallback for small real matrices."""
    shape = shape_of(matrix)
    if len(shape) > 2:
        return [eigenvalues_nested(item) for item in matrix]
    rows = [[float(value) for value in row] for row in as_2d(matrix)]
    if len(rows) != len(rows[0]):
        raise ValueError(f"Eigenvalues require a square matrix, got {shape}.")
    if len(rows) == 0:
        return []
    if len(rows) == 1:
        return [rows[0][0]]
    if len(rows) == 2:
        a, b = rows[0]
        c, d = rows[1]
        trace = a + d
        determinant = a * d - b * c
        discriminant = trace * trace - 4.0 * determinant
        if discriminant >= 0:
            root = math.sqrt(discriminant)
            return [(trace - root) / 2.0, (trace + root) / 2.0]
        real = trace / 2.0
        imaginary = math.sqrt(-discriminant) / 2.0
        return [complex(real, -imaginary), complex(real, imaginary)]
    if _is_symmetric_matrix(rows):
        values, _vectors = _symmetric_eigh_desc(rows)
        return values
    raise RuntimeError(
        "Pure-Python eigenvalue fallback supports 1x1, 2x2, and symmetric matrices. "
        "Install numpy or torch for general eigendecomposition."
    )

factored_frobenius_norm(left, right)

Return ||left @ right||_F from small Gram products.

Source code in src/SafeLens/core/factored_matrix.py 
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
def factored_frobenius_norm(left: Any, right: Any) -> Any:
    """Return `||left @ right||_F` from small Gram products."""
    try:
        import torch

        if isinstance(left, torch.Tensor) or isinstance(right, torch.Tensor):
            if not isinstance(left, torch.Tensor):
                left = torch.as_tensor(left, dtype=right.dtype, device=right.device)
            if not isinstance(right, torch.Tensor):
                right = torch.as_tensor(right, dtype=left.dtype, device=left.device)
            left_gram = transpose(left) @ left
            right_gram = right @ transpose(right)
            return torch.sqrt((left_gram * transpose(right_gram)).sum(dim=(-2, -1)))
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(left, "shape") or hasattr(right, "shape"):
            left_array = np.asarray(left)
            right_array = np.asarray(right)
            left_gram = np.swapaxes(left_array, -1, -2) @ left_array
            right_gram = right_array @ np.swapaxes(right_array, -1, -2)
            return np.sqrt(
                (left_gram * np.swapaxes(right_gram, -1, -2)).sum(axis=(-2, -1))
            ).tolist()
    except Exception:
        pass
    left_shape = shape_of(left)
    right_shape = shape_of(right)
    if len(left_shape) > 2 or len(right_shape) > 2:
        leading_shape = broadcast_leading_shape(left_shape[:-2], right_shape[:-2])
        left_b = broadcast_to_shape(left, leading_shape + left_shape[-2:])
        right_b = broadcast_to_shape(right, leading_shape + right_shape[-2:])
        return [
            factored_frobenius_norm(left_item, right_item)
            for left_item, right_item in zip(left_b, right_b, strict=True)
        ]
    left_gram = matmul(transpose(left), left)
    right_gram = matmul(right, transpose(right))
    return math.sqrt(
        sum(
            float(left_gram[row][col]) * float(right_gram[col][row])
            for row in range(len(left_gram))
            for col in range(len(left_gram[row]))
        )
    )

frobenius_norm(value)

Return Frobenius norm per leading dimension.

Source code in src/SafeLens/core/factored_matrix.py 
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
def frobenius_norm(value: Any) -> Any:
    """Return Frobenius norm per leading dimension."""
    try:
        import torch

        if isinstance(value, torch.Tensor):
            return torch.linalg.matrix_norm(value, ord="fro", dim=(-2, -1))
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(value, "shape"):
            array = np.asarray(value)
            return np.linalg.norm(array, axis=(-2, -1)).tolist()
    except Exception:
        pass
    shape = shape_of(value)
    if len(shape) <= 2:
        return math.sqrt(sum(float(item) ** 2 for item in flatten_values(value)))
    return [frobenius_norm(item) for item in value]

matmul(left, right)

Matrix multiply tensor-like or nested-list values.

Source code in src/SafeLens/core/factored_matrix.py 
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def matmul(left: Any, right: Any) -> Any:
    """Matrix multiply tensor-like or nested-list values."""
    try:
        return left @ right
    except Exception:
        pass
    backend_result = _try_backend_matmul(left, right)
    if backend_result is not _BACKEND_MATMUL_NOT_AVAILABLE:
        return backend_result
    left_shape = shape_of(left)
    right_shape = shape_of(right)
    if len(left_shape) > 2 and len(right_shape) > 2:
        leading_shape = broadcast_leading_shape(left_shape[:-2], right_shape[:-2])
        left_b = broadcast_to_shape(left, leading_shape + left_shape[-2:])
        right_b = broadcast_to_shape(right, leading_shape + right_shape[-2:])
        return [
            matmul(left_item, right_item)
            for left_item, right_item in zip(left_b, right_b, strict=True)
        ]
    if len(left_shape) > 2 and len(right_shape) == 1:
        return [matmul(item, right) for item in left]
    if len(left_shape) == 1 and len(right_shape) > 2:
        return [matmul(left, item) for item in right]
    if len(left_shape) > 2 and len(right_shape) == 2:
        return [matmul(item, right) for item in left]
    if len(left_shape) == 2 and len(right_shape) > 2:
        left_b = broadcast_to_shape(left, right_shape[:-2] + left_shape)
        return [
            matmul(left_item, right_item)
            for left_item, right_item in zip(left_b, right, strict=True)
        ]
    if len(left_shape) == 1 and len(right_shape) == 1:
        return sum(float(a) * float(b) for a, b in zip(left, right, strict=True))
    if len(left_shape) == 1 and len(right_shape) == 2:
        right_t = transpose(right)
        return [sum(float(a) * float(b) for a, b in zip(left, col, strict=True)) for col in right_t]
    left_rows = as_2d(left)
    right_rows = as_2d(right)
    if len(right_rows) == 1 and len(left_rows[0]) != 1:
        vector = right_rows[0]
        return [
            sum(float(a) * float(b) for a, b in zip(row, vector, strict=True)) for row in left_rows
        ]
    right_t = transpose(right_rows)
    return [
        [sum(float(a) * float(b) for a, b in zip(row, col, strict=True)) for col in right_t]
        for row in left_rows
    ]

matrix_corner(matrix, k=3)

Return a top-left corner over the final two matrix dimensions.

Source code in src/SafeLens/core/factored_matrix.py 
714
715
716
717
718
719
720
721
722
723
724
725
726
def matrix_corner(matrix: Any, k: int = 3) -> Any:
    """Return a top-left corner over the final two matrix dimensions."""
    shape = shape_of(matrix)
    if len(shape) < 2:
        return matrix
    try:
        return matrix[..., :k, :k]
    except Exception:
        pass
    if len(shape) > 2:
        return [matrix_corner(item, k) for item in matrix]
    rows = as_2d(matrix)
    return [row[:k] for row in rows[:k]]

multiply_values(left, right)

Multiply tensor-like or nested values elementwise.

Source code in src/SafeLens/core/factored_matrix.py 
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
def multiply_values(left: Any, right: Any) -> Any:
    """Multiply tensor-like or nested values elementwise."""
    try:
        return left * right
    except Exception:
        pass
    backend_result = _try_backend_elementwise(left, right, op="multiply")
    if backend_result is not _BACKEND_MATMUL_NOT_AVAILABLE:
        return backend_result
    left_shape = shape_of(left)
    right_shape = shape_of(right)
    if left_shape or right_shape:
        left_b, right_b = broadcast_pair(left, right)
        if _is_sequence(left_b) and _is_sequence(right_b):
            return [
                multiply_values(left_item, right_item)
                for left_item, right_item in zip(left_b, right_b, strict=True)
            ]
    return float(left) * float(right)

scale_columns(matrix, column_scales)

Multiply each column of a matrix-like value by the matching scale.

Source code in src/SafeLens/core/factored_matrix.py 
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
def scale_columns(matrix: Any, column_scales: Any) -> Any:
    """Multiply each column of a matrix-like value by the matching scale."""
    try:
        return matrix * column_scales[..., None, :]
    except Exception:
        pass
    matrix_shape = shape_of(matrix)
    scale_shape = shape_of(column_scales)
    if len(matrix_shape) > 2 or len(scale_shape) > 1:
        return [
            scale_columns(matrix_item, scale_item)
            for matrix_item, scale_item in zip(matrix, column_scales, strict=True)
        ]
    rows = as_2d(matrix)
    scales = _as_1d(column_scales)
    return [
        [float(value) * float(scales[col_index]) for col_index, value in enumerate(row)]
        for row in rows
    ]

scale_rows(matrix, row_scales)

Multiply each row of a matrix-like value by the matching scale.

Source code in src/SafeLens/core/factored_matrix.py 
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
def scale_rows(matrix: Any, row_scales: Any) -> Any:
    """Multiply each row of a matrix-like value by the matching scale."""
    try:
        return row_scales[..., :, None] * matrix
    except Exception:
        pass
    matrix_shape = shape_of(matrix)
    scale_shape = shape_of(row_scales)
    if len(matrix_shape) > 2 or len(scale_shape) > 1:
        return [
            scale_rows(matrix_item, scale_item)
            for matrix_item, scale_item in zip(matrix, row_scales, strict=True)
        ]
    rows = as_2d(matrix)
    scales = _as_1d(row_scales)
    return [
        [float(value) * float(scales[row_index]) for value in row]
        for row_index, row in enumerate(rows)
    ]

scale_values(value, scalar)

Scale tensor-like or nested-list values by a scalar.

Source code in src/SafeLens/core/factored_matrix.py 
450
451
452
453
454
455
456
457
458
def scale_values(value: Any, scalar: Any) -> Any:
    """Scale tensor-like or nested-list values by a scalar."""
    if _is_sequence(value):
        return [scale_values(item, scalar) for item in value]
    try:
        return value * scalar
    except TypeError:
        pass
    return float(value) * float(scalar)

shape_of(value)

Return best-effort shape for a tensor-like value.

Source code in src/SafeLens/core/factored_matrix.py 
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
def shape_of(value: Any) -> tuple[int, ...]:
    """Return best-effort shape for a tensor-like value."""
    shape = getattr(value, "shape", None)
    if shape is not None:
        return tuple(int(dim) for dim in shape)
    if _is_sequence(value):
        if not value:
            return (0,)
        return (len(value), *shape_of(value[0]))
    return ()

squeeze_dim(value, dim=None)

Squeeze tensor-like or nested-list values.

Source code in src/SafeLens/core/factored_matrix.py 
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
def squeeze_dim(value: Any, dim: int | None = None) -> Any:
    """Squeeze tensor-like or nested-list values."""
    squeeze = getattr(value, "squeeze", None)
    if callable(squeeze):
        return squeeze() if dim is None else squeeze(dim)
    shape = shape_of(value)
    if dim is None:
        result = value
        for axis in reversed([index for index, size in enumerate(shape) if size == 1]):
            result = squeeze_dim(result, axis)
        return result
    if dim < 0:
        dim = len(shape) + dim
    if not shape or shape[dim] != 1:
        raise ValueError(f"Cannot squeeze dimension {dim} with shape {shape}.")
    return squeeze_nested_dim(value, dim)

svd_nested(matrix)

Small pure-Python SVD fallback for dependency-light list backends.

Source code in src/SafeLens/core/factored_matrix.py 
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
def svd_nested(matrix: Any) -> tuple[Any, Any, Any]:
    """Small pure-Python SVD fallback for dependency-light list backends."""
    shape = shape_of(matrix)
    if len(shape) > 2:
        u_items: list[Any] = []
        s_items: list[Any] = []
        v_items: list[Any] = []
        for item in matrix:
            u_item, s_item, v_item = svd_nested(item)
            u_items.append(u_item)
            s_items.append(s_item)
            v_items.append(v_item)
        return u_items, s_items, v_items
    rows = [[float(value) for value in row] for row in as_2d(matrix)]
    if not rows:
        return [], [], []
    row_count = len(rows)
    column_count = len(rows[0])
    if row_count == 0 or column_count == 0:
        return [[] for _ in range(row_count)], [], [[] for _ in range(column_count)]
    diagonal = _diagonal_values_if_diagonal(rows)
    if diagonal is not None:
        return _diagonal_svd(rows, diagonal)
    return _svd_via_symmetric_eigen(rows)

transpose(matrix)

Transpose the final two dimensions of a matrix-like value.

Source code in src/SafeLens/core/factored_matrix.py 
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
def transpose(matrix: Any) -> Any:
    """Transpose the final two dimensions of a matrix-like value."""
    shape = shape_of(matrix)
    if len(shape) >= 2 and not _is_sequence(matrix):
        matrix_transpose = getattr(matrix, "mT", None)
        if matrix_transpose is not None:
            return matrix_transpose
        swapaxes = getattr(matrix, "swapaxes", None)
        if callable(swapaxes):
            try:
                return swapaxes(-1, -2)
            except Exception:
                pass
        transpose_fn = getattr(matrix, "transpose", None)
        if callable(transpose_fn):
            try:
                return transpose_fn(-1, -2)
            except Exception:
                pass
        transpose_attr = getattr(matrix, "T", None)
        if transpose_attr is not None and len(shape) == 2:
            return transpose_attr
    if len(shape) > 2 and _is_sequence(matrix):
        return [transpose(item) for item in matrix]
    rows = as_2d(matrix)
    return [list(col) for col in zip(*rows, strict=True)]

unsqueeze_dim(value, dim)

Unsqueeze tensor-like or nested-list values.

Source code in src/SafeLens/core/factored_matrix.py 
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
def unsqueeze_dim(value: Any, dim: int) -> Any:
    """Unsqueeze tensor-like or nested-list values."""
    unsqueeze = getattr(value, "unsqueeze", None)
    if callable(unsqueeze):
        return unsqueeze(dim)
    shape = shape_of(value)
    rank = len(shape)
    if dim < 0:
        dim = rank + dim + 1
    if dim <= 0:
        return [value]
    if _is_sequence(value):
        return [unsqueeze_dim(item, dim - 1) for item in value]
    return [value]