Skip to content

Operator reference

Base

The :class:LinearOperator base class.

LinearOperator

Abstract base class for matrix-free finite-dimensional linear operators.

It follows in most parts the implementation of probnum.linops.LinearOperator from ProbNum (see also https://github.com/probabilistic-numerics/probnum/blob/main/src/probnum/linops/_linear_operator.py).

Design choices: :class:LinearOperator\ s are defined to behave like a :class:jax.numpy.ndarray and thus, they

  • have :attr:shape, :attr:dtype, :attr:ndim, and :attr:size attributes,
  • can be matrix multiplied (:code:@) with a :class:numpy.ndarray from left and right, following the same broadcasting rules as :func:numpy.matmul,
  • can be multiplied (:code:*) by a scalar from the left and the right,
  • can be added to, subtracted from and matrix multiplied (:code:@) with other :class:LinearOperator instances with appropriate :attr:shape,
  • can be transposed (:attr:T or :meth:transpose), and they
  • can be type-cast (:meth:astype).

This is mostly implemented lazily, i.e. the result of these operations is a new, composite :class:LinearOperator, that defers linear operations to the original operators and combines the results.

Parameters:

Name Type Description Default
shape ShapeLike

Shape of the linear operator.

 required
dtype DTypeLike
required
See Also

aslinop : Transform into a LinearOperator

Notes
  • A subclass is only required to implement :meth:_matmat. Additionally, other methods like :meth:_solve, :meth:_inverse, :meth:_transpose, :meth:_cholesky, or :meth:_det should be overwritten if more performant implementations are available.
  • Compared to probnum this implementation does not check for dtype to be numeric and not complexfloating.
  • Matrix properties are tags.
Important:
  • (...batch..., n, m) is the general shape assumption.
Source code in linox/operators/base.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
312
313
314
class LinearOperator:
    r"""Abstract base class for `matrix-free` finite-dimensional linear operators.

    It follows in most parts the implementation of `probnum.linops.LinearOperator`
    from ProbNum
    (see also https://github.com/probabilistic-numerics/probnum/blob/main/src/probnum/linops/_linear_operator.py).

    Design choices:
    :class:`LinearOperator`\ s are defined to behave like a :class:`jax.numpy.ndarray`
    and thus, they


    * have :attr:`shape`, :attr:`dtype`, :attr:`ndim`, and :attr:`size` attributes,
    * can be matrix multiplied (:code:`@`) with a :class:`numpy.ndarray` from left and
      right, following the same broadcasting rules as :func:`numpy.matmul`,
    * can be multiplied (:code:`*`) by a scalar from the left and the right,
    * can be added to, subtracted from and matrix multiplied (:code:`@`) with other
      :class:`LinearOperator` instances with appropriate :attr:`shape`,
    * can be transposed (:attr:`T` or :meth:`transpose`), and they
    * can be type-cast (:meth:`astype`).

    This is mostly implemented lazily, i.e. the result of these operations is a new,
    composite :class:`LinearOperator`, that defers linear operations to the original
    operators and combines the results.

    Parameters
    ----------
    shape: Tuple[int]
        Shape of the linear operator.
    dtype: Type

    See Also
    --------
    aslinop : Transform into a LinearOperator

    Notes
    -----
    -   A subclass is only required to implement :meth:`_matmat`. Additionally, other
        methods like :meth:`_solve`, :meth:`_inverse`, :meth:`_transpose`,
        :meth:`_cholesky`, or :meth:`_det` should be overwritten if more performant
        implementations are available.
    -   Compared to probnum this implementation does not check for dtype to be numeric
    and not complexfloating.
    -   Matrix properties are tags.

    Important:
    ----------
    -  (...batch..., n, m) is the general shape assumption.
    """

    def __init__(
        self,
        shape: ShapeLike,
        dtype: DTypeLike,
    ) -> None:
        self.__shape = utils.as_shape(shape, ndim=len(shape))

        self.__dtype = jnp.dtype(dtype)

    @property
    def shape(self) -> tuple[int]:
        """Shape of the linear operator.

        Defined as a tuple of the output and input dimension of operator.
        """
        return self.__shape

    @property
    def batch_shape(self) -> tuple[int]:
        """Shape of the batch dimensions of the linear operator."""
        return self.__shape[:-2]

    @property
    def ndim(self) -> int:
        """Number of linear operator dimensions.

        Defined analogously to numpy.ndarray.ndim.
        TODO(2bys): Check with jnp.ndarray.ndim.
        """
        return len(self.__shape)

    @property
    def batch_ndim(self) -> int:
        """Number of batch dimensions of the linear operator."""
        return len(self.__shape[:-2])

    @property
    def size(self) -> int:
        """Product of the :attr:`shape` entries."""
        return reduce(operator.mul, self.__shape, 1)

    @property
    def dtype(self) -> jnp.dtype:
        """Data type of the linear operator."""
        return self.__dtype

    @property
    def is_symmetric(self) -> bool:
        """Whether the operator is symmetric."""
        return False

    @property
    def is_psd(self) -> bool:
        """Whether the operator is positive semi-definite."""
        return False

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__} with shape={self.shape}, dtype={self.dtype}>"

    def graph(self, **kwargs):
        """Return this operator's structure as a tree of :class:`LinOpNode`."""
        from linox.utils.debug import linop_graph

        return linop_graph(self, **kwargs)

    def graph_str(self, **kwargs):
        """Return this operator's structure as a printable tree."""
        return self.graph(**kwargs).pretty()

    ########################################################################
    # Default Methods that should be overwritten
    ########################################################################

    def todense(self) -> jnp.ndarray:
        """Materialize this operator as a dense array."""
        if config.get_warn_on_densify():
            config.warn(f"Linear operator {self} is densed.", prefix="PerformanceWarning")

        config.emit(
            config.DebugEvent(
                kind="densify",
                msg="todense() called",
                op_type=type(self).__name__,
                op_id=id(self),
                shape=getattr(self, "shape", None),
                dtype=getattr(self, "dtype", None),
            )
        )
        return self @ jnp.eye(self.shape[-1], dtype=self.dtype)

    def _todense(self) -> jnp.ndarray:
        msg = "Subclasses must implement _todense"
        raise NotImplementedError(msg)

    def _matmul(self, other: jnp.ndarray) -> jnp.ndarray:
        return self._todense() @ other

    def transpose(self) -> "LinearOperator":
        """Return the transpose of this operator.

        Subclasses that know their own structure should override this and
        return it (``Diagonal`` returns itself, ``Kronecker`` returns a
        ``Kronecker`` of transposed factors, and so on). The default is a lazy
        wrapper that derives the adjoint from the forward matvec, so it never
        materialises the dense matrix.
        """
        from linox.operators.arithmetic import (
            TransposedLinearOperator,
        )

        return TransposedLinearOperator(self)

    @property
    def T(self) -> "LinearOperator":
        """Return the transpose of this operator, preserving structure where possible."""
        from linox.operators.arithmetic import (
            TransposedLinearOperator,
        )

        # Prefer a subclass's structured transpose (e.g. Diagonal -> itself,
        # Sym/PSD -> itself) so the operator's structure survives `.T`. The
        # base implementation returns a dense array rather than an operator,
        # in which case fall back to the lazy wrapper.
        transposed = self.transpose()
        if isinstance(transposed, LinearOperator):
            return transposed
        return TransposedLinearOperator(self)

    ########################################################################
    # Arithmetic
    ########################################################################

    def __neg__(self) -> "LinearOperator":
        from .arithmetic import lneg

        return lneg(self)

    def __add__(self, other: "LinearOperator") -> "LinearOperator":  # Here the package uses a BinaryOperandType
        from .arithmetic import ladd

        return ladd(self, other)

    def __radd__(self, other: BinaryOperandType) -> "LinearOperator":
        from .arithmetic import ladd

        # Addition is commutative, so reuse the forward dispatch: it carries
        # the (LinearOperator, Array) methods that the reversed argument order
        # would not resolve.
        return ladd(self, other)

    def __sub__(self, other: BinaryOperandType) -> "LinearOperator":
        from .arithmetic import lsub

        return lsub(self, other)

    def __rsub__(self, other: BinaryOperandType) -> "LinearOperator":
        from .arithmetic import lsub

        return lsub(other, self)

    def __mul__(self, other: BinaryOperandType) -> "LinearOperator":
        from .arithmetic import lmul

        return lmul(self, other)

    def __rmul__(self, other: BinaryOperandType) -> "LinearOperator":
        from .arithmetic import lmul

        return lmul(other, self)

    def __truediv__(self, other: BinaryOperandType) -> "LinearOperator":
        from .arithmetic import ldiv

        return ldiv(self, other)

    def __matmul__(self, other: BinaryOperandType) -> "LinearOperator":
        from .arithmetic import lmatmul

        flatten = False
        operand = other
        if isinstance(other, (jax.Array, np.ndarray)):
            operand = jnp.asarray(other)
            if operand.ndim == 1:
                operand = operand[:, None]
                flatten = True

        res = lmatmul(self, operand)
        if (
            not flatten
            and isinstance(res, jax.Array)
            and res.ndim >= 2
            and res.shape[-2] != self.shape[-2]
            and hasattr(operand, "shape")
            and res.shape[-2] == operand.shape[-1]
            and res.shape[-1] == self.shape[-2]
        ):
            res = jnp.swapaxes(res, -1, -2)
        return res if not flatten else res[..., 0]

    def __rmatmul__(self, other: BinaryOperandType) -> "LinearOperator":
        from linox.operators.arithmetic import lmatmul

        # lazy evaluation
        isLazyEvaluation = True

        if other.shape[-1] != self.shape[-2]:
            msg = f"expected other.shape[-1] to be {other.shape[-1]}, got {self.shape[-2]} instead."
            raise ValueError(msg)

        if len(other.shape) > 2:
            msg = "Only 2D arrays are supported."
            raise ValueError(msg)

        if len(other.shape) == 1:
            other = other[None, :]
            isLazyEvaluation = False

        res = lmatmul(other, self)
        return res if isLazyEvaluation else (res[0, :] if isinstance(res, jnp.ndarray) else res._todense()[0])

    def __call__(self, arr: BinaryOperandType) -> "LinearOperator":
        """Apply this operator, equivalent to ``self @ arr``."""
        return self @ arr

    @classmethod
    def tree_flatten(cls) -> tuple[tuple[any, ...], dict[str, any]]:
        """Default implementation for PyTree flattening.

        Subclasses should override this method to provide proper PyTree support.
        """
        children = ()  # No children by default
        aux_data = {}  # No auxiliary data by default
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "LinearOperator":
        """Default implementation for PyTree unflattening."""
        del children
        if cls is LinearOperator:
            msg = "Cannot unflatten the abstract LinearOperator class directly"
            raise TypeError(msg)
        return cls(**aux_data)

T: LinearOperator property

Return the transpose of this operator, preserving structure where possible.

batch_ndim: int property

Number of batch dimensions of the linear operator.

batch_shape: tuple[int] property

Shape of the batch dimensions of the linear operator.

dtype: jnp.dtype property

Data type of the linear operator.

is_psd: bool property

Whether the operator is positive semi-definite.

is_symmetric: bool property

Whether the operator is symmetric.

ndim: int property

Number of linear operator dimensions.

Defined analogously to numpy.ndarray.ndim. TODO(2bys): Check with jnp.ndarray.ndim.

shape: tuple[int] property

Shape of the linear operator.

Defined as a tuple of the output and input dimension of operator.

size: int property

Product of the :attr:shape entries.

__call__(arr: BinaryOperandType) -> LinearOperator

Apply this operator, equivalent to self @ arr.

Source code in linox/operators/base.py 
289
290
291
def __call__(self, arr: BinaryOperandType) -> "LinearOperator":
    """Apply this operator, equivalent to ``self @ arr``."""
    return self @ arr

graph(**kwargs)

Return this operator's structure as a tree of :class:LinOpNode.

Source code in linox/operators/base.py 
128
129
130
131
132
def graph(self, **kwargs):
    """Return this operator's structure as a tree of :class:`LinOpNode`."""
    from linox.utils.debug import linop_graph

    return linop_graph(self, **kwargs)

graph_str(**kwargs)

Return this operator's structure as a printable tree.

Source code in linox/operators/base.py 
134
135
136
def graph_str(self, **kwargs):
    """Return this operator's structure as a printable tree."""
    return self.graph(**kwargs).pretty()

todense() -> jnp.ndarray

Materialize this operator as a dense array.

Source code in linox/operators/base.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
def todense(self) -> jnp.ndarray:
    """Materialize this operator as a dense array."""
    if config.get_warn_on_densify():
        config.warn(f"Linear operator {self} is densed.", prefix="PerformanceWarning")

    config.emit(
        config.DebugEvent(
            kind="densify",
            msg="todense() called",
            op_type=type(self).__name__,
            op_id=id(self),
            shape=getattr(self, "shape", None),
            dtype=getattr(self, "dtype", None),
        )
    )
    return self @ jnp.eye(self.shape[-1], dtype=self.dtype)

transpose() -> LinearOperator

Return the transpose of this operator.

Subclasses that know their own structure should override this and return it (Diagonal returns itself, Kronecker returns a Kronecker of transposed factors, and so on). The default is a lazy wrapper that derives the adjoint from the forward matvec, so it never materialises the dense matrix.

Source code in linox/operators/base.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
def transpose(self) -> "LinearOperator":
    """Return the transpose of this operator.

    Subclasses that know their own structure should override this and
    return it (``Diagonal`` returns itself, ``Kronecker`` returns a
    ``Kronecker`` of transposed factors, and so on). The default is a lazy
    wrapper that derives the adjoint from the forward matvec, so it never
    materialises the dense matrix.
    """
    from linox.operators.arithmetic import (
        TransposedLinearOperator,
    )

    return TransposedLinearOperator(self)

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]] classmethod

Default implementation for PyTree flattening.

Subclasses should override this method to provide proper PyTree support.

Source code in linox/operators/base.py 
293
294
295
296
297
298
299
300
301
@classmethod
def tree_flatten(cls) -> tuple[tuple[any, ...], dict[str, any]]:
    """Default implementation for PyTree flattening.

    Subclasses should override this method to provide proper PyTree support.
    """
    children = ()  # No children by default
    aux_data = {}  # No auxiliary data by default
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> LinearOperator classmethod

Default implementation for PyTree unflattening.

Source code in linox/operators/base.py 
303
304
305
306
307
308
309
310
311
312
313
314
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "LinearOperator":
    """Default implementation for PyTree unflattening."""
    del children
    if cls is LinearOperator:
        msg = "Cannot unflatten the abstract LinearOperator class directly"
        raise TypeError(msg)
    return cls(**aux_data)

Basic

Dense Matrix operator.

Matrix

Bases: LinearOperator

A linear operator defined via a matrix.

For a matrix :math:A, this represents the linear operator :math:x \mapsto Ax. The action on a vector :math:x is given by matrix multiplication :math:Ax.

Args: A: The matrix defining the linear operator

Source code in linox/operators/dense.py 
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
class Matrix(LinearOperator):
    r"""A linear operator defined via a matrix.

    For a matrix :math:`A`, this represents the linear operator :math:`x \mapsto Ax`.
    The action on a vector :math:`x` is given by matrix multiplication :math:`Ax`.

    Args:
        A: The matrix defining the linear operator
    """

    def __init__(self, A: ArrayLike) -> None:  # type: ignore
        self.A = jnp.asarray(A)
        config.emit(
            config.DebugEvent(
                kind="init",
                msg=f"Matrix initialized with shape {self.A.shape} and dtype {self.A.dtype}",
                op_type=type(self).__name__,
                op_id=id(self),
                shape=self.A.shape,
                dtype=self.A.dtype,
            )
        )
        super().__init__(self.A.shape, self.A.dtype)

    def _matmul(self, vector: jax.Array) -> jax.Array:
        return self.A @ vector

    def _todense(self) -> jax.Array:
        _warn(f"Converting Matrix of shape {self.shape} to dense array.")
        return self.A

    def transpose(self) -> "Matrix":
        """Return the transpose of this operator."""
        return Matrix(self.A.swapaxes(-1, -2))

    def __T__(self) -> "Matrix":
        """Alias for transpose."""
        return self.transpose()

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = (self.A,)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data: dict[str, any], children: tuple[any, ...]) -> "Matrix":
        """Reconstruct this operator from JAX pytree children and static data."""
        del aux_data
        (A,) = children
        return cls(A=A)

__T__() -> Matrix

Alias for transpose.

Source code in linox/operators/dense.py 
58
59
60
def __T__(self) -> "Matrix":
    """Alias for transpose."""
    return self.transpose()

transpose() -> Matrix

Return the transpose of this operator.

Source code in linox/operators/dense.py 
54
55
56
def transpose(self) -> "Matrix":
    """Return the transpose of this operator."""
    return Matrix(self.A.swapaxes(-1, -2))

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/dense.py 
62
63
64
65
66
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = (self.A,)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> Matrix classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/dense.py 
68
69
70
71
72
73
@classmethod
def tree_unflatten(cls, aux_data: dict[str, any], children: tuple[any, ...]) -> "Matrix":
    """Reconstruct this operator from JAX pytree children and static data."""
    del aux_data
    (A,) = children
    return cls(A=A)

Diagonal Linear Operators.

Diagonal

Bases: LinearOperator

A linear operator defined via a diagonal matrix.

For a vector :math:d, this represents the diagonal matrix :math:\text{diag}(d). The action on a vector :math:x is given by element-wise multiplication :math:\text{diag}(d)x = d \odot x where :math:\odot denotes element-wise multiplication.

Args: diag: The diagonal elements of the matrix

Source code in linox/operators/diagonal.py 
 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
class Diagonal(LinearOperator):
    r"""A linear operator defined via a diagonal matrix.

    For a vector :math:`d`, this represents the diagonal matrix :math:`\text{diag}(d)`.
    The action on a vector :math:`x` is given by element-wise multiplication
    :math:`\text{diag}(d)x = d \odot x` where :math:`\odot` denotes element-wise
    multiplication.

    Args:
        diag: The diagonal elements of the matrix
    """

    def __init__(self, diag: ArrayLike) -> None:
        if isinstance(diag, Diagonal):
            diag = diag.diag
        self.diag = jnp.asarray(diag)
        # Read the shape off the converted array, not the raw argument, so
        # list/tuple inputs work.
        super().__init__(
            shape=(
                *self.diag.shape[:-1],
                self.diag.shape[-1],
                self.diag.shape[-1],
            ),
            dtype=self.diag.dtype,
        )

    def _matmul(self, vector: jax.Array) -> jax.Array:
        return self.diag[..., None] * vector

    def _todense(self) -> jax.Array:
        return _batch_jnp_diag(self.diag)

    @property
    def is_symmetric(self) -> bool:
        # Real diagonal matrices are symmetric.
        # If complex, check imaginary part?
        # Generally yes for simplicity in real case.
        """Whether this operator equals its own transpose (always true here)."""
        return True

    @property
    def is_psd(self) -> bool:
        """Whether every diagonal entry is non-negative.

        Returns ``False`` for a traced diagonal rather than guessing: under
        ``jax.jit`` the entries are tracers with no concrete ordering, so no
        claim can be made. A bare ``except`` here previously swallowed every
        error, including unrelated bugs.
        """
        try:
            return bool(jnp.all(self.diag >= 0))
        except jax.errors.ConcretizationTypeError:
            return False

    def transpose(self) -> "Diagonal":
        """Return the transpose of this operator."""
        return self

    def diagonal(self) -> jax.Array:
        """Return the diagonal entries."""
        return self.diag

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = (self.diag,)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "Diagonal":
        """Reconstruct this operator from JAX pytree children and static data."""
        del aux_data
        (diag,) = children
        return cls(diag=diag)

is_psd: bool property

Whether every diagonal entry is non-negative.

Returns False for a traced diagonal rather than guessing: under jax.jit the entries are tracers with no concrete ordering, so no claim can be made. A bare except here previously swallowed every error, including unrelated bugs.

is_symmetric: bool property

Whether this operator equals its own transpose (always true here).

diagonal() -> jax.Array

Return the diagonal entries.

Source code in linox/operators/diagonal.py 
87
88
89
def diagonal(self) -> jax.Array:
    """Return the diagonal entries."""
    return self.diag

transpose() -> Diagonal

Return the transpose of this operator.

Source code in linox/operators/diagonal.py 
83
84
85
def transpose(self) -> "Diagonal":
    """Return the transpose of this operator."""
    return self

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/diagonal.py 
91
92
93
94
95
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = (self.diag,)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> Diagonal classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/diagonal.py 
 97
 98
 99
100
101
102
103
104
105
106
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "Diagonal":
    """Reconstruct this operator from JAX pytree children and static data."""
    del aux_data
    (diag,) = children
    return cls(diag=diag)

Special Linear Operators (Identity, Zero, Scalar, Ones).

Identity

Bases: LinearOperator

The identity operator.

This represents the identity matrix :math:I. The action on a vector :math:x is given by :math:Ix = x, i.e., the identity operator leaves vectors unchanged.

Args: shape: The shape of the identity operator dtype: The data type of the identity operator. Defaults to JAX's current floating dtype (float64 under x64, else float32).

Source code in linox/operators/special.py 
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
class Identity(LinearOperator):
    r"""The identity operator.

    This represents the identity matrix :math:`I`. The action on a vector :math:`x` is
    given by :math:`Ix = x`, i.e., the identity operator leaves vectors unchanged.

    Args:
        shape: The shape of the identity operator
        dtype: The data type of the identity operator. Defaults to JAX's
            current floating dtype (float64 under x64, else float32).
    """

    def __init__(self, shape: ShapeLike, *, dtype: DTypeLike | None = None) -> None:
        dtype = default_floating_dtype() if dtype is None else dtype
        shape = as_shape(shape)
        super().__init__((*shape, shape[-1]), dtype)

    def _matmul(self, arr: jax.Array) -> jax.Array:
        return jnp.broadcast_to(
            arr,
            shape=(
                *jnp.broadcast_shapes(arr.shape[:-2], self.shape[:-2]),
                self.shape[-2],
                arr.shape[-1],
            ),
        )

    def _todense(self) -> jax.Array:
        return jnp.broadcast_to(jnp.eye(self.shape[-1], dtype=self.dtype), self.shape)

    @property
    def is_symmetric(self) -> bool:
        """Check if operator is symmetric."""
        return True

    @property
    def is_psd(self) -> bool:
        """Check if operator is positive semi-definite."""
        return True

    def transpose(self) -> "Identity":
        """Return transpose (self for identity)."""
        return self

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration."""
        children = ()
        aux_data = {"shape": self.shape[:-1], "dtype": self.dtype}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "Identity":
        """Unflatten for JAX pytree registration."""
        del children
        return cls(shape=aux_data["shape"], dtype=aux_data["dtype"])

is_psd: bool property

Check if operator is positive semi-definite.

is_symmetric: bool property

Check if operator is symmetric.

transpose() -> Identity

Return transpose (self for identity).

Source code in linox/operators/special.py 
80
81
82
def transpose(self) -> "Identity":
    """Return transpose (self for identity)."""
    return self

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

Source code in linox/operators/special.py 
84
85
86
87
88
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration."""
    children = ()
    aux_data = {"shape": self.shape[:-1], "dtype": self.dtype}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> Identity classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/special.py 
90
91
92
93
94
95
96
97
98
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "Identity":
    """Unflatten for JAX pytree registration."""
    del children
    return cls(shape=aux_data["shape"], dtype=aux_data["dtype"])

Ones

Bases: LinearOperator

The ones operator.

This represents the matrix :math:\mathbf{1}\mathbf{1}^T where :math:\mathbf{1} is a vector of ones. The action on a vector :math:x is given by :math:(\mathbf{1}\mathbf{1}^T)x = \mathbf{1}(\mathbf{1}^T x), i.e., it sums the elements of :math:x and returns a vector of that sum.

Args: shape: The shape of the ones operator dtype: The data type of the ones operator. Defaults to JAX's current floating dtype (float64 under x64, else float32).

Source code in linox/operators/special.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
class Ones(LinearOperator):
    r"""The ones operator.

    This represents the matrix :math:`\mathbf{1}\mathbf{1}^T` where :math:`\mathbf{1}`
    is a vector of ones. The action on a vector :math:`x` is given by
    :math:`(\mathbf{1}\mathbf{1}^T)x = \mathbf{1}(\mathbf{1}^T x)`, i.e., it sums the
    elements of :math:`x` and returns a vector of that sum.

    Args:
        shape: The shape of the ones operator
        dtype: The data type of the ones operator. Defaults to JAX's
            current floating dtype (float64 under x64, else float32).
    """

    def __init__(self, shape: ShapeLike, dtype: DTypeLike | None = None) -> None:
        dtype = default_floating_dtype() if dtype is None else dtype
        super().__init__(shape, dtype)

    def _matmul(self, arr: jax.Array) -> jax.Array:
        return jnp.broadcast_to(
            arr.sum(axis=-2, keepdims=True),
            shape=(
                *jnp.broadcast_shapes(arr.shape[:-2], self.shape[:-2]),
                self.shape[-2],
                arr.shape[-1],
            ),
        )

    def _todense(self) -> jax.Array:
        return jnp.ones(self.shape, dtype=self.dtype)

    def transpose(self) -> "Ones":
        """Return transposed ones operator."""
        return Ones(shape=(*self.shape[:-2], self.shape[-1], self.shape[-2]), dtype=self.dtype)

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration."""
        children = ()
        aux_data = {"shape": self.shape, "dtype": self.dtype}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "Ones":
        """Unflatten for JAX pytree registration."""
        del children
        return cls(shape=aux_data["shape"], dtype=aux_data["dtype"])

transpose() -> Ones

Return transposed ones operator.

Source code in linox/operators/special.py 
523
524
525
def transpose(self) -> "Ones":
    """Return transposed ones operator."""
    return Ones(shape=(*self.shape[:-2], self.shape[-1], self.shape[-2]), dtype=self.dtype)

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

Source code in linox/operators/special.py 
527
528
529
530
531
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration."""
    children = ()
    aux_data = {"shape": self.shape, "dtype": self.dtype}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> Ones classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/special.py 
533
534
535
536
537
538
539
540
541
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "Ones":
    """Unflatten for JAX pytree registration."""
    del children
    return cls(shape=aux_data["shape"], dtype=aux_data["dtype"])

Scalar

Bases: LinearOperator

A linear operator defined via a scalar.

For a scalar :math:\alpha, this represents :math:\alpha I where :math:I is the identity matrix. The action on a vector :math:x is given by scalar multiplication :math:(\alpha I)x = \alpha x.

Args: scalar: The scalar value defining the operator

Source code in linox/operators/special.py 
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
class Scalar(LinearOperator):
    r"""A linear operator defined via a scalar.

    For a scalar :math:`\alpha`, this represents :math:`\alpha I` where :math:`I`
    is the identity matrix. The action on a vector :math:`x` is given by scalar
    multiplication :math:`(\alpha I)x = \alpha x`.

    Args:
        scalar: The scalar value defining the operator
    """

    def __init__(self, scalar: ScalarLike) -> None:
        self.scalar = jnp.asarray(scalar)

        super().__init__(shape=(), dtype=self.scalar.dtype)

    def _matmul(self, vector: jax.Array) -> jax.Array:
        return self.scalar * vector

    def _todense(self) -> jax.Array:
        return self.scalar  # Scalar todense returns the scalar array? Or should it expand?
        # _matrix.py said `return self`. Wait, self is the operator. returning self in _todense is weird unless it means the scalar value?
        # Ah, lines 622 in _matrix.py: `return self`. This looks like a bug in original code or self.scalar?
        # Wait, if I return self (the LinearOperator instance), that's definitely wrong for `todense`.
        # However, line 622 in previous output says `return self`.
        # Wait, if `Scalar` acts like a scalar array, maybe?
        # But `todense()` is expected to return jax.Array.
        # Let's fix it to return `self.scalar`.
        return self.scalar

    @property
    def is_symmetric(self) -> bool:
        """Check if operator is symmetric."""
        return True

    @property
    def is_psd(self) -> bool:
        """Check if operator is positive semi-definite."""
        try:
            return float(self.scalar) >= 0
        except Exception:
            return False

    def transpose(self) -> "Scalar":
        """Return transpose (self for scalar)."""
        return self

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration."""
        children = (self.scalar,)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "Scalar":
        """Unflatten for JAX pytree registration."""
        del aux_data
        (scalar,) = children
        return cls(scalar=scalar)

is_psd: bool property

Check if operator is positive semi-definite.

is_symmetric: bool property

Check if operator is symmetric.

transpose() -> Scalar

Return transpose (self for scalar).

Source code in linox/operators/special.py 
291
292
293
def transpose(self) -> "Scalar":
    """Return transpose (self for scalar)."""
    return self

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

Source code in linox/operators/special.py 
295
296
297
298
299
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration."""
    children = (self.scalar,)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> Scalar classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/special.py 
301
302
303
304
305
306
307
308
309
310
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "Scalar":
    """Unflatten for JAX pytree registration."""
    del aux_data
    (scalar,) = children
    return cls(scalar=scalar)

Zero

Bases: LinearOperator

The zero operator.

This represents the zero matrix :math:0. The action on a vector :math:x is given by :math:0x = 0, i.e., the zero operator maps all vectors to zero.

Args: shape: The shape of the zero operator dtype: The data type of the zero operator. Defaults to JAX's current floating dtype (float64 under x64, else float32).

Source code in linox/operators/special.py 
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
class Zero(LinearOperator):
    r"""The zero operator.

    This represents the zero matrix :math:`0`. The action on a vector :math:`x` is
    given by :math:`0x = 0`, i.e., the zero operator maps all vectors to zero.

    Args:
        shape: The shape of the zero operator
        dtype: The data type of the zero operator. Defaults to JAX's
            current floating dtype (float64 under x64, else float32).
    """

    def __init__(self, shape: ShapeLike, dtype: DTypeLike | None = None) -> None:
        dtype = default_floating_dtype() if dtype is None else dtype
        super().__init__(shape, dtype)

    def _matmul(self, arr: jax.Array) -> jax.Array:
        return jnp.zeros(
            (
                *jnp.broadcast_shapes(arr.shape[:-2], self.shape[:-2]),
                self.shape[-2],
                arr.shape[-1],
            ),
            # Promote with the operand: a Zero operator must not narrow the
            # result of `Zero @ x` to its own dtype.
            dtype=jnp.result_type(self.dtype, arr.dtype),
        )

    def _todense(self) -> jax.Array:
        return jnp.zeros(self.shape, dtype=self.dtype)

    @property
    def is_symmetric(self) -> bool:
        """Check if operator is symmetric."""
        return True

    @property
    def is_psd(self) -> bool:
        """Check if operator is positive semi-definite."""
        return True

    def transpose(self) -> "Zero":
        """Return transposed zero operator."""
        return Zero(shape=(*self.shape[:-2], self.shape[-1], self.shape[-2]), dtype=self.dtype)

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration."""
        children = ()
        aux_data = {"shape": self.shape, "dtype": self.dtype}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "Zero":
        """Unflatten for JAX pytree registration."""
        del children
        return cls(shape=aux_data["shape"], dtype=aux_data["dtype"])

is_psd: bool property

Check if operator is positive semi-definite.

is_symmetric: bool property

Check if operator is symmetric.

transpose() -> Zero

Return transposed zero operator.

Source code in linox/operators/special.py 
408
409
410
def transpose(self) -> "Zero":
    """Return transposed zero operator."""
    return Zero(shape=(*self.shape[:-2], self.shape[-1], self.shape[-2]), dtype=self.dtype)

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

Source code in linox/operators/special.py 
412
413
414
415
416
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration."""
    children = ()
    aux_data = {"shape": self.shape, "dtype": self.dtype}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> Zero classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/special.py 
418
419
420
421
422
423
424
425
426
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "Zero":
    """Unflatten for JAX pytree registration."""
    del children
    return cls(shape=aux_data["shape"], dtype=aux_data["dtype"])

Composition

Arithmetic operations for linear operators.

This module implements various arithmetic operations for linear operators, including:

  • :class:ScaledLinearOperator: Represents :math:\alpha A for scalar :math:\alpha and operator :math:A
  • :class:AddLinearOperator: Represents :math:A_1 + A_2 + \ldots + A_n for operators :math:A_i
  • :class:ProductLinearOperator: Represents :math:A_1A_2\ldots A_n for operators :math:A_i
  • :class:CongruenceTransform: Represents :math:ABA^T for operators :math:A and :math:B
  • :class:TransposedLinearOperator: Represents :math:A^T for operator :math:A
  • :class:InverseLinearOperator: Represents :math:A^{-1} for operator :math:A

These operators can be combined to form complex linear transformations while maintaining efficient computation through lazy evaluation.

ScaledLinearOperator

Bases: LinearOperator

Linear operator scaled with a scalar.

For a linear operator :math:A and scalar :math:\alpha, this represents :math:\alpha A where :math:(\alpha A)x = \alpha(Ax) for any vector :math:x

Args: operator: A linear operator to be scaled scalar: A scalar value to multiply the operator with

Source code in linox/operators/arithmetic.py 
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
class ScaledLinearOperator(LinearOperator):
    r"""Linear operator scaled with a scalar.

    For a linear operator :math:`A` and scalar :math:`\alpha`, this represents
    :math:`\alpha A` where :math:`(\alpha A)x = \alpha(Ax)` for any vector :math:`x`

    Args:
        operator: A linear operator to be scaled
        scalar: A scalar value to multiply the operator with
    """

    def __init__(self, operator: LinearOperator, scalar: ScalarLike) -> None:
        self.operator = utils.as_linop(operator)
        scalar = jnp.asarray(scalar)
        dtype = jnp.result_type(operator.dtype, scalar.dtype)
        self.scalar = utils.as_scalar(scalar, dtype)
        super().__init__(shape=operator.shape, dtype=dtype)

    @property
    def is_symmetric(self) -> bool:
        """Whether the scaled operator is symmetric."""
        return self.operator.is_symmetric

    @property
    def is_psd(self) -> bool:
        # Check if scalar >= 0 and operator is PSD
        # Note: robust checking of scalar value requires jax.Array value access?
        # If scalar is tracer, this might fail or return tracer.
        # Introspection is usually for planning (outside JIT), so concrete values expected.
        # But here scalar is stored as array/scalar.
        # For safety, we can try converting to float if it's not traced.
        # A traced scalar has no concrete sign, so make no claim rather than
        # guessing. (A bare `except` here previously swallowed every error.)
        """Whether the scaled operator is positive semi-definite."""
        try:
            s = float(self.scalar)
        except (jax.errors.ConcretizationTypeError, TypeError):
            return False
        return s >= 0 and self.operator.is_psd

    def _matmul(self, arr: jax.Array) -> jax.Array:
        return (self.operator @ arr) * self.scalar

    def _todense(self) -> jax.Array:
        return self.operator._todense() * self.scalar

    def transpose(self) -> LinearOperator:
        """Return the transpose of this operator."""
        return self.scalar * self.operator.transpose()

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = (self.operator, self.scalar)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "ScaledLinearOperator":
        """Reconstruct this operator from JAX pytree children and static data."""
        del aux_data
        operator, scalar = children
        return cls(operator=operator, scalar=scalar)

is_psd: bool property

Whether the scaled operator is positive semi-definite.

is_symmetric: bool property

Whether the scaled operator is symmetric.

transpose() -> LinearOperator

Return the transpose of this operator.

Source code in linox/operators/arithmetic.py 
1010
1011
1012
def transpose(self) -> LinearOperator:
    """Return the transpose of this operator."""
    return self.scalar * self.operator.transpose()

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1014
1015
1016
1017
1018
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = (self.operator, self.scalar)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> ScaledLinearOperator classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "ScaledLinearOperator":
    """Reconstruct this operator from JAX pytree children and static data."""
    del aux_data
    operator, scalar = children
    return cls(operator=operator, scalar=scalar)

AddLinearOperator

Bases: LinearOperator

A linear operator formed by adding two or more linear operators together.

For linear operators :math:A_1, A_2, \ldots, A_n, this represents :math:A_1 + A_2 + \ldots + A_nwhere :math:(A_1 + A_2 + \ldots + A_n)x = A_1x + A_2x + \ldots + A_nx for any vector :math:x

Args: *operator_list: Variable number of linear operators to be added

Source code in linox/operators/arithmetic.py 
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
class AddLinearOperator(LinearOperator):
    r"""A linear operator formed by adding two or more linear operators together.

    For linear operators :math:`A_1, A_2, \ldots, A_n`, this represents
    :math:`A_1 + A_2 + \ldots + A_n`where
    :math:`(A_1 + A_2 + \ldots + A_n)x = A_1x + A_2x + \ldots + A_nx`
    for any vector :math:`x`

    Args:
        *operator_list: Variable number of linear operators to be added
    """

    def __init__(self, *operator_list: ArithmeticType) -> None:
        self.operator_list = [
            utils.as_linop(o) if isinstance(op, AddLinearOperator) else utils.as_linop(op)
            for op in operator_list
            for o in (op.operator_list if isinstance(op, AddLinearOperator) else [op])
        ]
        shape = _broadcast_shapes([op.shape for op in self.operator_list])
        super().__init__(shape=shape, dtype=self.operator_list[0].dtype)

    def _matmul(self, arr: jax.Array) -> jax.Array:
        return reduce(
            operator.add,
            (op @ arr for op in reversed(self.operator_list)),
        )

    def _todense(self) -> jax.Array:
        return reduce(operator.add, (op._todense() for op in self.operator_list))

    def transpose(self) -> "AddLinearOperator":
        """Return the transpose of this operator."""
        return AddLinearOperator(*(op.transpose() for op in self.operator_list))

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = tuple(self.operator_list)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "AddLinearOperator":
        """Reconstruct this operator from JAX pytree children and static data."""
        del aux_data
        return cls(*children)

transpose() -> AddLinearOperator

Return the transpose of this operator.

Source code in linox/operators/arithmetic.py 
1134
1135
1136
def transpose(self) -> "AddLinearOperator":
    """Return the transpose of this operator."""
    return AddLinearOperator(*(op.transpose() for op in self.operator_list))

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1138
1139
1140
1141
1142
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = tuple(self.operator_list)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> AddLinearOperator classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1144
1145
1146
1147
1148
1149
1150
1151
1152
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "AddLinearOperator":
    """Reconstruct this operator from JAX pytree children and static data."""
    del aux_data
    return cls(*children)

ProductLinearOperator

Bases: LinearOperator

Product of linear operators.

For linear operators :math:A_1, A_2, \ldots, A_n, this represents :math:A_1A_2\ldots A_n where :math:(A_1A_2\ldots A_n)x = A_1(A_2(\ldots(A_nx))) for any vector :math:x

Args: *operator_list: Variable number of linear operators to be multiplied

Source code in linox/operators/arithmetic.py 
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
class ProductLinearOperator(LinearOperator):
    r"""Product of linear operators.

    For linear operators :math:`A_1, A_2, \ldots, A_n`, this represents
    :math:`A_1A_2\ldots A_n` where :math:`(A_1A_2\ldots A_n)x = A_1(A_2(\ldots(A_nx)))`
    for any vector :math:`x`

    Args:
        *operator_list: Variable number of linear operators to be multiplied
    """

    def __init__(self, *operator_list: LinearOperator) -> None:
        self.operator_list = [
            utils.as_linop(o) if isinstance(op, ProductLinearOperator) else utils.as_linop(op)
            for op in operator_list
            for o in (op.operator_list if isinstance(op, ProductLinearOperator) else [op])
        ]
        batch_shape = _broadcast_shapes([op.shape[:-2] for op in self.operator_list])
        self.__check_init__()
        result_dtype = jnp.result_type(*[op.dtype for op in self.operator_list])
        shape = utils.as_shape(
            (
                *batch_shape,
                self.operator_list[0].shape[-2],
                self.operator_list[-1].shape[-1],
            )
        )
        super().__init__(shape=shape, dtype=result_dtype)

    def __check_init__(self) -> None:
        for i, op1 in enumerate(self.operator_list[:-1]):
            op2 = self.operator_list[i + 1]
            if op1.shape[-1] != op2.shape[-2]:
                msg = f"Shape mismatch: Cannot multiply linear operators with shapes operator 1: ({op1.shape}) operator 2: ({op2.shape})"
                raise ValueError(msg)

    def _matmul(self, arr: jax.Array) -> jax.Array:
        return reduce(lambda x, y: y @ x, [arr, *reversed(self.operator_list)])

    def transpose(self) -> "ProductLinearOperator":
        """Return the transpose of this operator."""
        return ProductLinearOperator(*(op.transpose() for op in reversed(self.operator_list)))

    def _todense(self) -> jax.Array:
        return reduce(
            lambda x, y: y @ x,
            [
                self.operator_list[-1]._todense(),
                *reversed([op._todense() for op in self.operator_list[:-1]]),
            ],
        )

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = tuple(self.operator_list)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data: dict[str, any], children: tuple[any, ...]) -> "ProductLinearOperator":
        """Reconstruct this operator from JAX pytree children and static data."""
        del aux_data
        return cls(*children)

transpose() -> ProductLinearOperator

Return the transpose of this operator.

Source code in linox/operators/arithmetic.py 
1239
1240
1241
def transpose(self) -> "ProductLinearOperator":
    """Return the transpose of this operator."""
    return ProductLinearOperator(*(op.transpose() for op in reversed(self.operator_list)))

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1252
1253
1254
1255
1256
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = tuple(self.operator_list)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> ProductLinearOperator classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1258
1259
1260
1261
1262
@classmethod
def tree_unflatten(cls, aux_data: dict[str, any], children: tuple[any, ...]) -> "ProductLinearOperator":
    """Reconstruct this operator from JAX pytree children and static data."""
    del aux_data
    return cls(*children)

TransposedLinearOperator

Bases: LinearOperator

Transpose of a linear operator.

For a linear operator :math:A, this represents :math:A^T where :math:(A^T)_{ij} = A_{ji} for all :math:i,j

Args: operator: A linear operator to be transposed

Source code in linox/operators/arithmetic.py 
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
class TransposedLinearOperator(LinearOperator):
    r"""Transpose of a linear operator.

    For a linear operator :math:`A`, this represents :math:`A^T`
    where :math:`(A^T)_{ij} = A_{ji}` for all :math:`i,j`

    Args:
        operator: A linear operator to be transposed
    """

    def __init__(self, operator: LinearOperator) -> None:
        self.operator = utils.as_linop(operator)
        batch_shape = operator.shape[:-2]
        super().__init__(
            shape=(*batch_shape, operator.shape[-1], operator.shape[-2]),
            dtype=operator.dtype,
        )

    def _matmul(self, arr: jnp.array) -> jax.Array:
        # Derive the adjoint from the forward matvec instead of asking the
        # wrapped operator to transpose itself. `LinearOperator.transpose`
        # materialises the dense matrix by default, so the old
        # `self.operator.transpose() @ arr` densified any operator that did
        # not override it -- exactly the operators a matrix-free library must
        # not densify.
        #
        # This wrapper is only ever constructed when there is *no* structured
        # transpose (`LinearOperator.T` returns the subclass's own transpose
        # when it provides one), so nothing structured loses its fast path.
        arr = jnp.asarray(arr)
        n = self.operator.shape[-1]

        # `linear_transpose` requires the cotangent's dtype to match the
        # forward output's exactly, so promote both to a common dtype rather
        # than silently narrowing a float64 rhs onto a float32 operator.
        dtype = jnp.promote_types(self.operator.dtype, arr.dtype)
        basis = jnp.zeros((n,), dtype=dtype)

        def forward(v: jax.Array) -> jax.Array:
            return self.operator @ v

        out_dtype = jax.eval_shape(forward, basis).dtype

        def adjoint(col: jax.Array) -> jax.Array:
            return jax.linear_transpose(forward, basis)(col.astype(out_dtype))[0]

        if arr.ndim == 1:
            return adjoint(arr)
        return jax.vmap(adjoint, in_axes=-1, out_axes=-1)(arr)

    def _todense(self) -> jax.Array:
        return self.operator._todense().swapaxes(-1, -2)

    def transpose(self) -> LinearOperator:
        """Return the transpose of this operator."""
        return self.operator

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = (self.operator,)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "TransposedLinearOperator":
        """Reconstruct this operator from JAX pytree children and static data."""
        del aux_data
        (operator,) = children
        return cls(operator=operator)

transpose() -> LinearOperator

Return the transpose of this operator.

Source code in linox/operators/arithmetic.py 
1371
1372
1373
def transpose(self) -> LinearOperator:
    """Return the transpose of this operator."""
    return self.operator

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1375
1376
1377
1378
1379
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = (self.operator,)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> TransposedLinearOperator classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "TransposedLinearOperator":
    """Reconstruct this operator from JAX pytree children and static data."""
    del aux_data
    (operator,) = children
    return cls(operator=operator)

InverseLinearOperator

Bases: LinearOperator

Inverse of a linear operator.

For a linear operator :math:A, this represents :math:A^{-1} where :math:A^{-1} is the unique operator such that :math:AA^{-1} = A^{-1}A = I where :math:I is the identity operator

Args: operator: A linear operator to be inverted

Source code in linox/operators/arithmetic.py 
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
class InverseLinearOperator(LinearOperator):
    """Inverse of a linear operator.

    For a linear operator :math:`A`, this represents :math:`A^{-1}`
    where :math:`A^{-1}` is the unique operator such that :math:`AA^{-1} = A^{-1}A = I`
    where :math:`I` is the identity operator

    Args:
        operator: A linear operator to be inverted
    """

    def __init__(
        self,
        operator: LinearOperator,
        method: str = "exact",
        solver_options: dict[str, any] | None = None,
    ) -> None:
        self.operator = operator
        self.method = method
        self.solver_options = solver_options if solver_options is not None else {}
        super().__init__(shape=operator.shape, dtype=operator.dtype)

    def _matmul(self, arr: jax.Array) -> jax.Array:
        if self.method == "exact":
            return lsolve(self.operator, arr)

        if self.method == "lsmr":
            from linox.linalg.approx.lsmr import lsmr_solve

            x, _ = lsmr_solve(self.operator, arr, **self.solver_options)
            return x

        if self.method in {"cg", "conjugate_gradient"}:
            x, _ = jax.scipy.sparse.linalg.cg(self.operator, arr, **self.solver_options)
            return x

        # Fallback
        return lsolve(self.operator, arr)

    def _todense(self) -> jax.Array:
        if self.method == "exact":
            config.warn(f"Linear operator {self.operator} is densed for inverse computation.")
            return jnp.linalg.inv(self.operator._todense())
        # If approx, we can't easily densify without solving against identity
        # Fallback to solving against Identity (expensive but correct)
        n = self.shape[0]
        I_op = jnp.eye(n, dtype=self.dtype)
        return self._matmul(I_op)

    def transpose(self) -> LinearOperator:
        """Return the transpose of this operator."""
        # If A is invertible, (A^-1)^T = (A^T)^-1
        # We propagate the solver method.
        # Note: if method was "cg" (expects SPD), A^T should also be SPD if A was.
        # If method was "lsmr", it works for A^T too.
        return InverseLinearOperator(
            self.operator.transpose(),
            method=self.method,
            solver_options=self.solver_options,
        )

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = (self.operator,)

        aux_data = {"method": self.method, "solver_options": self.solver_options} if self.method != "exact" else {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "InverseLinearOperator":
        """Reconstruct this operator from JAX pytree children and static data."""
        (operator,) = children
        return cls(operator=operator, **aux_data)

transpose() -> LinearOperator

Return the transpose of this operator.

Source code in linox/operators/arithmetic.py 
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
def transpose(self) -> LinearOperator:
    """Return the transpose of this operator."""
    # If A is invertible, (A^-1)^T = (A^T)^-1
    # We propagate the solver method.
    # Note: if method was "cg" (expects SPD), A^T should also be SPD if A was.
    # If method was "lsmr", it works for A^T too.
    return InverseLinearOperator(
        self.operator.transpose(),
        method=self.method,
        solver_options=self.solver_options,
    )

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1460
1461
1462
1463
1464
1465
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = (self.operator,)

    aux_data = {"method": self.method, "solver_options": self.solver_options} if self.method != "exact" else {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> InverseLinearOperator classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1467
1468
1469
1470
1471
1472
1473
1474
1475
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "InverseLinearOperator":
    """Reconstruct this operator from JAX pytree children and static data."""
    (operator,) = children
    return cls(operator=operator, **aux_data)

PseudoInverseLinearOperator

Bases: LinearOperator

Moore-Penrose pseudo-inverse A^+ of a linear operator.

Source code in linox/operators/arithmetic.py 
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
class PseudoInverseLinearOperator(LinearOperator):
    """Moore-Penrose pseudo-inverse ``A^+`` of a linear operator."""

    def __init__(self, operator: LinearOperator, tol: float = 1e-12) -> None:
        self.operator = operator
        super().__init__(shape=operator.T.shape, dtype=operator.dtype)
        self.tol = tol

    def transpose(self) -> LinearOperator:
        """Return the transpose of this operator."""
        # (A^+)^T == (A^T)^+ -- transpose the operand, not the pseudo-inverse
        # of self, which recurses forever.
        return PseudoInverseLinearOperator(self.operator.transpose(), tol=self.tol)

    def _todense(self) -> jax.Array:
        r"""Materialize the pseudo-inverse densely.

        TODO: compute this from the SVD rather than via ``jnp.linalg.pinv``:
        ``U, S, Vh = svd(self.operator)``.

        Returns
        -------
            x_LS = \sum_i (u_i^T b) / s_i v_i
            -> U, S, Vh = svd(self.operator)
            return U @ jnp.diag(1 / S) @ Vh.
        """
        return jnp.linalg.pinv(self.operator._todense(), rtol=self.tol)

    def _matmul(self, arr: jax.Array) -> jax.Array:
        return lpsolve(self.operator, arr, rtol=self.tol)

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = (self.operator,)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "PseudoInverseLinearOperator":
        """Reconstruct this operator from JAX pytree children and static data."""
        del aux_data
        (operator,) = children
        return cls(operator=operator)

transpose() -> LinearOperator

Return the transpose of this operator.

Source code in linox/operators/arithmetic.py 
1544
1545
1546
1547
1548
def transpose(self) -> LinearOperator:
    """Return the transpose of this operator."""
    # (A^+)^T == (A^T)^+ -- transpose the operand, not the pseudo-inverse
    # of self, which recurses forever.
    return PseudoInverseLinearOperator(self.operator.transpose(), tol=self.tol)

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1567
1568
1569
1570
1571
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = (self.operator,)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> PseudoInverseLinearOperator classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "PseudoInverseLinearOperator":
    """Reconstruct this operator from JAX pytree children and static data."""
    del aux_data
    (operator,) = children
    return cls(operator=operator)

CongruenceTransform

Bases: ProductLinearOperator

:math:A B A^\top.

Source code in linox/operators/arithmetic.py 
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
class CongruenceTransform(ProductLinearOperator):
    r""":math:`A B A^\top`."""

    def __init__(self, A: ArithmeticType, B: ArithmeticType) -> None:
        self._A = utils.as_linop(A)
        self._B = utils.as_linop(B)

        super().__init__(self._A, self._B, self._A.T)

    def transpose(self) -> LinearOperator:
        """Return the transpose of this operator."""
        return CongruenceTransform(self._A, self._B.T)

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = (self._A, self._B)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "CongruenceTransform":
        """Reconstruct this operator from JAX pytree children and static data."""
        del aux_data
        A, B = children
        return cls(A=A, B=B)

transpose() -> LinearOperator

Return the transpose of this operator.

Source code in linox/operators/arithmetic.py 
1508
1509
1510
def transpose(self) -> LinearOperator:
    """Return the transpose of this operator."""
    return CongruenceTransform(self._A, self._B.T)

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1512
1513
1514
1515
1516
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = (self._A, self._B)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> CongruenceTransform classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/arithmetic.py 
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "CongruenceTransform":
    """Reconstruct this operator from JAX pytree children and static data."""
    del aux_data
    A, B = children
    return cls(A=A, B=B)

Structured

Kronecker product operations for linear operators.

This module includes:

  • :class:Kronecker: Represents the Kronecker product :math:A \otimes B of two linear operators :math:A and :math:B
  • :class:KroneckerSelectedEigenvectors: Matrix-free operator for selected eigenvectors of a Kronecker product
  • :func:topk_eigh: Compute top-k or bottom-k eigenvalues/vectors of a Kronecker product without forming the full matrix

Kronecker

Bases: LinearOperator

A Kronecker product of two linear operators.

Example usage:

A = jnp.array([[1, 2], [3, 4]], dtype=jnp.float32) B = jnp.array([[5, 6], [7, 8]], dtype=jnp.float32) op = Kronecker(A, B) vec = jnp.ones((4,)) result = op @ vec result_true = jnp.kron(A, B) @ vec jnp.allclose(result, result_true)

Source code in linox/operators/kron.py 
 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
class Kronecker(LinearOperator):
    """A Kronecker product of two linear operators.

    Example usage:

    A = jnp.array([[1, 2], [3, 4]], dtype=jnp.float32)
    B = jnp.array([[5, 6], [7, 8]], dtype=jnp.float32)
    op = Kronecker(A, B)
    vec = jnp.ones((4,))
    result = op @ vec
    result_true = jnp.kron(A, B) @ vec
    jnp.allclose(result, result_true)
    """

    def __init__(self, A: LinearOperator | jax.Array, B: LinearOperator | jax.Array) -> None:
        self._A = utils.as_linop(A)
        self._B = utils.as_linop(B)
        A_shape = self._A.shape if len(self._A.shape) == 2 else (self._A.shape[0], 1)
        B_shape = self._B.shape if len(self._B.shape) == 2 else (self._B.shape[0], 1)

        self._shape = (
            A_shape[0] * B_shape[0],
            A_shape[1] * B_shape[1],
        )

        dtype = jnp.result_type(self._A.dtype, self._B.dtype)
        super().__init__(self._shape, dtype)

    @property
    def A(self) -> LinearOperator:
        """First factor of the Kronecker product."""
        return self._A

    @property
    def B(self) -> LinearOperator:
        """Second factor of the Kronecker product."""
        return self._B

    @property
    def shape(self) -> tuple[int, int]:
        """Shape of the Kronecker product."""
        return self._shape

    @property
    def is_symmetric(self) -> bool:
        """Check if Kronecker product is symmetric."""
        return self.A.is_symmetric and self.B.is_symmetric

    @property
    def is_psd(self) -> bool:
        """Check if Kronecker product is positive semi-definite."""
        return self.A.is_psd and self.B.is_psd

    def tree_flatten(self) -> tuple[tuple, dict]:
        """Flatten for JAX pytree registration."""
        children = (self.A, self.B)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict,
        children: tuple,
    ) -> "Kronecker":
        """Unflatten for JAX pytree registration."""
        return cls(*children)

    def _matmul(self, vec: jax.Array) -> jax.Array:
        if len(vec.shape) == 1:
            vec = vec[:, None]

        _, mA = self.A.shape
        _, mB = self.B.shape

        y = jnp.swapaxes(vec, -2, -1)
        y = y.reshape((*y.shape[:-1], mA, mB))
        y = self.B @ jnp.swapaxes(y, -1, -2)
        y = self.A @ jnp.swapaxes(y, -1, -2)
        y = y.reshape((*y.shape[:-2], -1))
        y = jnp.swapaxes(y, -1, -2)

        return y

    def _todense(self) -> jax.Array:
        return jnp.kron(self.A._todense(), self.B._todense())

    def transpose(self) -> "Kronecker":
        """Return transposed Kronecker product."""
        return Kronecker(self.A.transpose(), self.B.transpose())

    def trace(self) -> jax.Array:
        """Compute trace of Kronecker product: tr(A (x) B) = tr(A) tr(B)."""
        # `self.A` / `self.B` are LinearOperators, so take their diagonals
        # through the dispatch rather than calling `jnp.trace` on them.
        return jnp.sum(jnp.asarray(diagonal(self.A)), axis=-1) * jnp.sum(jnp.asarray(diagonal(self.B)), axis=-1)

A: LinearOperator property

First factor of the Kronecker product.

B: LinearOperator property

Second factor of the Kronecker product.

is_psd: bool property

Check if Kronecker product is positive semi-definite.

is_symmetric: bool property

Check if Kronecker product is symmetric.

shape: tuple[int, int] property

Shape of the Kronecker product.

trace() -> jax.Array

Compute trace of Kronecker product: tr(A (x) B) = tr(A) tr(B).

Source code in linox/operators/kron.py 
137
138
139
140
141
def trace(self) -> jax.Array:
    """Compute trace of Kronecker product: tr(A (x) B) = tr(A) tr(B)."""
    # `self.A` / `self.B` are LinearOperators, so take their diagonals
    # through the dispatch rather than calling `jnp.trace` on them.
    return jnp.sum(jnp.asarray(diagonal(self.A)), axis=-1) * jnp.sum(jnp.asarray(diagonal(self.B)), axis=-1)

transpose() -> Kronecker

Return transposed Kronecker product.

Source code in linox/operators/kron.py 
133
134
135
def transpose(self) -> "Kronecker":
    """Return transposed Kronecker product."""
    return Kronecker(self.A.transpose(), self.B.transpose())

tree_flatten() -> tuple[tuple, dict]

Flatten for JAX pytree registration.

Source code in linox/operators/kron.py 
 99
100
101
102
103
def tree_flatten(self) -> tuple[tuple, dict]:
    """Flatten for JAX pytree registration."""
    children = (self.A, self.B)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict, children: tuple) -> Kronecker classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/kron.py 
105
106
107
108
109
110
111
112
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict,
    children: tuple,
) -> "Kronecker":
    """Unflatten for JAX pytree registration."""
    return cls(*children)

KroneckerSelectedEigenvectors

Bases: LinearOperator

Matrix-free operator for selected Kronecker eigenvectors.

Represents :math:Q_k where columns are :math:q_A^{(i)} \otimes q_B^{(j)} \otimes \ldots for selected index tuples.

Never forms the full Kronecker product. For two factors A, B with eigenvectors :math:Q_A, Q_B:

.. math:: Q_k \alpha = \text{vec}((U_B \odot \alpha) U_A^T)

where :math:U_A = Q_A[:, \text{selected}_A], :math:U_B = Q_B[:, \text{selected}_B].

Args: factor_vecs: List of (n_i, n_i) eigenvector matrices from each factor selected_indices: List of k tuples specifying which eigenvector combinations sort_indices: Sorting permutations applied to each factor's eigenvalues

Source code in linox/operators/kron.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
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
class KroneckerSelectedEigenvectors(LinearOperator):
    r"""Matrix-free operator for selected Kronecker eigenvectors.

    Represents :math:`Q_k` where columns are
    :math:`q_A^{(i)} \otimes q_B^{(j)} \otimes \ldots` for selected index tuples.

    Never forms the full Kronecker product. For two factors A, B with
    eigenvectors :math:`Q_A, Q_B`:

    .. math::
        Q_k \alpha = \text{vec}((U_B \odot \alpha) U_A^T)

    where :math:`U_A = Q_A[:, \text{selected}_A]`, :math:`U_B = Q_B[:, \text{selected}_B]`.

    Args:
        factor_vecs: List of (n_i, n_i) eigenvector matrices from each factor
        selected_indices: List of k tuples specifying which eigenvector combinations
        sort_indices: Sorting permutations applied to each factor's eigenvalues
    """

    def __init__(
        self,
        factor_vecs: list[jax.Array],
        selected_indices: list[tuple[int, ...]],
        sort_indices: list[jax.Array],
    ) -> None:
        self._factor_vecs = factor_vecs
        self._selected_indices = selected_indices
        self._sort_indices = sort_indices

        self._d = len(factor_vecs)
        self._k = len(selected_indices)
        self._factor_dims = [Q.shape[0] for Q in factor_vecs]
        self._n_total = int(jnp.prod(jnp.array(self._factor_dims)))

        sel_np = np.asarray(selected_indices, dtype=np.int32)  # (k, d)

        # `factor_vecs` arrive with their columns ALREADY permuted into
        # eigenvalue-sorted order (see `topk_eigh`), and `selected_indices`
        # index into that same sorted order. Re-applying `sort_indices` here
        # would permute a second time and select the wrong columns entirely --
        # in practice the bottom of the spectrum instead of the top.
        # `sort_indices` is retained only so callers can recover the mapping
        # back to each factor's original eigenvector ordering.
        self._gathered = []
        for i in range(self._d):
            # idx_sorted: (k,) -- indices into the already-sorted columns
            idx_sorted = jnp.asarray(sel_np[:, i])
            self._gathered.append(self._factor_vecs[i][:, idx_sorted])

        dtype = factor_vecs[0].dtype
        super().__init__((self._n_total, self._k), dtype)

    @property
    def k(self) -> int:
        """Number of selected eigenvector combinations."""
        return self._k

    @property
    def num_factors(self) -> int:
        """Number of Kronecker factors."""
        return self._d

    @property
    def factor_dims(self) -> list[int]:
        """Dimensions of each Kronecker factor."""
        return self._factor_dims

    def tree_flatten(self) -> tuple[tuple, dict]:
        """Flatten for JAX pytree registration."""
        children = (
            tuple(self._factor_vecs),
            tuple(self._sort_indices),
        )
        aux_data = {
            "selected_indices": self._selected_indices,
        }
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data: dict, children: tuple) -> "KroneckerSelectedEigenvectors":
        """Unflatten for JAX pytree registration."""
        factor_vecs, sort_indices = children
        return cls(
            list(factor_vecs),
            aux_data["selected_indices"],
            list(sort_indices),
        )

    def _matmul(self, alpha: jax.Array) -> jax.Array:
        r"""Compute :math:`Q_k \alpha` without forming the Kronecker product."""
        squeeze = False
        if alpha.ndim == 1:
            alpha = alpha[:, None]
            squeeze = True

        if self._d == 2:
            UA, UB = self._gathered[0], self._gathered[1]
            nA, nB = self._factor_dims[0], self._factor_dims[1]
            Y = jnp.einsum("il,lb,jl->ijb", UA, alpha, UB)
            result = Y.reshape((nA * nB, -1))

            if squeeze:
                result = result.squeeze(-1)
            return result

        result = jnp.zeros((self._n_total, alpha.shape[1]), dtype=self.dtype)

        for l in range(self._k):
            vec_l = self._gathered[0][:, l]
            for i in range(1, self._d):
                vec_l = jnp.kron(vec_l, self._gathered[i][:, l])
            result += vec_l[:, None] * alpha[l, :]

        if squeeze:
            result = result.squeeze(-1)
        return result

    def _rmatmul(self, v: jax.Array) -> jax.Array:
        r"""Compute Q_k^T v."""
        n = self._n_total
        restore = None

        if v.ndim == 1:
            pass

        elif v.ndim == 2:
            if v.shape[0] == n and v.shape[1] != n:
                v = jnp.swapaxes(v, 0, 1)
                restore = ("cols",)

            elif v.shape[1] == n:
                restore = ("batch",)

            elif v.shape[1] == 1 and v.shape[0] == n:
                v = v[:, 0]
            else:
                msg = f"Unsupported v shape {v.shape}. Expected (n,), (batch,n), or (n,p)"
                raise ValueError(msg)
        else:
            msg = f"Unsupported v.ndim={v.ndim}. Expected 1 or 2."
            raise ValueError(msg)

        squeeze_single = False
        if v.ndim == 1:
            v = v[None, :]
            squeeze_single = True

        if self._d == 2:
            UA, UB = self._gathered[0], self._gathered[1]
            nA, nB = self._factor_dims[0], self._factor_dims[1]

            X = v.reshape((v.shape[0], nA, nB))
            T = jnp.einsum("il,bij->blj", UA, X)
            result = jnp.einsum("blj,jl->bl", T, UB)

        else:
            result = jnp.zeros((v.shape[0], self._k), dtype=self.dtype)
            for l in range(self._k):
                vec_l = self._gathered[0][:, l]
                for i in range(1, self._d):
                    vec_l = jnp.kron(vec_l, self._gathered[i][:, l])
                result = result.at[:, l].set(v @ vec_l)

        if squeeze_single:
            result = result[0, :]

        if restore == ("cols",) and result.ndim == 2:
            result = jnp.swapaxes(result, 0, 1)

        return result

    def transpose(self) -> "KroneckerSelectedEigenvectorsTranspose":
        """Return transpose operator."""
        return KroneckerSelectedEigenvectorsTranspose(self)

    def _todense(self) -> jax.Array:
        cols = []
        for l in range(self._k):
            vec_l = self._gathered[0][:, l]
            for i in range(1, self._d):
                vec_l = jnp.kron(vec_l, self._gathered[i][:, l])
            cols.append(vec_l)
        return jnp.stack(cols, axis=1)

factor_dims: list[int] property

Dimensions of each Kronecker factor.

k: int property

Number of selected eigenvector combinations.

num_factors: int property

Number of Kronecker factors.

transpose() -> KroneckerSelectedEigenvectorsTranspose

Return transpose operator.

Source code in linox/operators/kron.py 
664
665
666
def transpose(self) -> "KroneckerSelectedEigenvectorsTranspose":
    """Return transpose operator."""
    return KroneckerSelectedEigenvectorsTranspose(self)

tree_flatten() -> tuple[tuple, dict]

Flatten for JAX pytree registration.

Source code in linox/operators/kron.py 
560
561
562
563
564
565
566
567
568
569
def tree_flatten(self) -> tuple[tuple, dict]:
    """Flatten for JAX pytree registration."""
    children = (
        tuple(self._factor_vecs),
        tuple(self._sort_indices),
    )
    aux_data = {
        "selected_indices": self._selected_indices,
    }
    return children, aux_data

tree_unflatten(aux_data: dict, children: tuple) -> KroneckerSelectedEigenvectors classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/kron.py 
571
572
573
574
575
576
577
578
579
@classmethod
def tree_unflatten(cls, aux_data: dict, children: tuple) -> "KroneckerSelectedEigenvectors":
    """Unflatten for JAX pytree registration."""
    factor_vecs, sort_indices = children
    return cls(
        list(factor_vecs),
        aux_data["selected_indices"],
        list(sort_indices),
    )

topk_eigh(op_or_factors, k: int, *, largest: bool = True, sigma2: float | jax.Array | None = None, include_noise_shift: bool = False, return_full_eigs: bool = False, mode: str = 'jax')

Compute top-k eigenvalues/vectors of a Kronecker product.

Parameters:

Name Type Description Default
op_or_factors LinearOperator or list[LinearOperator]

Either a Kronecker operator or a list of factor operators

 required
k int

Number of eigenvalues/vectors to compute

 required
largest bool

If True, compute largest eigenvalues; if False, compute smallest

 True
sigma2 float or Array

Noise variance for whitening

 None
include_noise_shift bool

Whether to include noise shift in eigenvalues

 False
return_full_eigs bool

If True, return full eigenvalue arrays for each factor

 False
mode str

Computation mode

 "jax"

Returns:

Name Type Description
eigenvalues Array

Top-k eigenvalues
eigenvectors LinearOperator

Top-k eigenvectors as a linear operator
info KronTopkEighInfo

Factorized representation with additional information
Source code in linox/operators/kron.py 
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
def topk_eigh(
    op_or_factors,
    k: int,
    *,
    largest: bool = True,
    sigma2: float | jax.Array | None = None,
    include_noise_shift: bool = False,
    return_full_eigs: bool = False,
    mode: str = "jax",
):
    """Compute top-k eigenvalues/vectors of a Kronecker product.

    Parameters
    ----------
    op_or_factors : LinearOperator or list[LinearOperator]
        Either a Kronecker operator or a list of factor operators
    k : int
        Number of eigenvalues/vectors to compute
    largest : bool, default=True
        If True, compute largest eigenvalues; if False, compute smallest
    sigma2 : float or jax.Array, optional
        Noise variance for whitening
    include_noise_shift : bool, default=False
        Whether to include noise shift in eigenvalues
    return_full_eigs : bool, default=False
        If True, return full eigenvalue arrays for each factor
    mode : str, default="jax"
        Computation mode

    Returns
    -------
    eigenvalues : jax.Array
        Top-k eigenvalues
    eigenvectors : LinearOperator
        Top-k eigenvectors as a linear operator
    info : KronTopkEighInfo
        Factorized representation with additional information
    """
    scalar = None
    if isinstance(op_or_factors, LinearOperator):
        factors, scalar = extract_kronecker_factors(op_or_factors)
    else:
        factors = [utils.as_linop(f) for f in op_or_factors]

    factor_eigs: list[jax.Array] = []
    factor_vecs: list[jax.Array] = []
    sort_indices: list[jax.Array] = []
    full_factor_eigs: list[jax.Array] = []

    # compute eigendecomp of each factor
    for A in factors:
        w, Q = leigh(A)
        if isinstance(w, LinearOperator):
            w = diagonal(w)

        eps = jnp.finfo(w.dtype).eps
        w_safe = jnp.maximum(w, eps)  # (n,)

        Q_dense = Q._todense() if hasattr(Q, "_todense") else jnp.asarray(Q)  # (n,n)

        order = jnp.argsort(-w_safe) if largest else jnp.argsort(w_safe)  # (n,)
        w_sorted = w_safe[order]
        Q_sorted = Q_dense[:, order]

        sort_indices.append(order)
        factor_eigs.append(w_sorted)
        factor_vecs.append(Q_sorted)

        if return_full_eigs:
            full_factor_eigs.append(w_sorted)

    # numerical eps cutoff (host float)
    dtype = factor_eigs[0].dtype

    # time_start = time.time()  # Timing removed

    if mode == "jax":
        # JAX branch: NO device_get, NO python if on traced values
        scalar_j = jnp.asarray(1.0, dtype=dtype) if scalar is None else jnp.asarray(scalar, dtype=dtype)
        add_shift_j = jnp.asarray(0.0, dtype=dtype)
        if include_noise_shift and sigma2 is not None:
            add_shift_j = jnp.asarray(sigma2, dtype=dtype)

        eig_array, selected_indices = _topk_product_grid_indices_jax(
            factor_eigs,
            k,
            largest=largest,
            scalar=scalar_j,
            add_shift=add_shift_j,
        )

    else:
        # Host/debug branch only
        eps_cutoff = float(np.finfo(np.dtype(dtype)).eps)

        scalar_f = 1.0
        if scalar is not None:
            scalar_f = float(jax.device_get(jnp.asarray(scalar)))
            if scalar_f < 0.0:
                raise ValueError("Negative scalar breaks PSD monotone-grid assumptions.")

        add_shift = 0.0
        if include_noise_shift and sigma2 is not None:
            add_shift = float(jax.device_get(jnp.asarray(sigma2)))

        w_host_list = [np.asarray(jax.device_get(w)) for w in factor_eigs]
        eigvals_list, selected_indices = _topk_product_grid_indices_host(
            w_host_list,
            k,
            largest=largest,
            eps_cutoff=eps_cutoff,
            scalar=scalar_f,
            add_shift=add_shift,
        )
        eig_array = jnp.asarray(eigvals_list, dtype=dtype)

    # time_end = time.time() # Timing removed
    # logger.info(...) # Removed

    Qk = KroneckerSelectedEigenvectors(factor_vecs, selected_indices, sort_indices)

    info = KronTopkEighInfo(
        factor_vecs=factor_vecs,
        factor_eigs=factor_eigs,
        sort_indices=sort_indices,
        selected_indices=selected_indices,
        scalar=scalar,
    )
    if return_full_eigs:
        # return eig_array, Qk, info, full_factor_eigs
        # Compatibility wrapper: existing calls might expect (vals, vecs).
        # But this function is new (renamed from topk_eigh which was slightly different).
        # We'll return (vals, vecs, info, ...) and update users.
        return eig_array, Qk, info, full_factor_eigs
    return eig_array, Qk, info

Isotropic additive operators of the form s*I + A.

IsotropicAdditiveLinearOperator

Bases: AddLinearOperator

Isotropic additive linear operator for matrices of the form.

A_iso := s I + A,

where s is a scalar (or a 0-arg scalar LinearOperator) and A is a symmetric LinearOperator. This class exposes fast, matrix-free implementations of common spectral transforms (inverse, pseudo-inverse, square root, log, powers, exp, Cholesky-like factor) by working in the eigenbasis of A.

Core idea

If A = Q Λ Qᵀ is an eigendecomposition of A (with Λ diagonal and Qᵀ Q = I), then

s I + A = Q (Λ + s I) Qᵀ,

so any spectral function f (e.g. inverse, sqrt, log, power, exp) satisfies

f(s I + A) = Q f(Λ + s I) Qᵀ,

which reduces the linear-algebra to elementwise operations on the eigenvalues.

This class computes/caches an (optionally truncated) eigendecomposition via leigh(A) and then dispatches the following:

  • linverse: (s I + A)⁻¹ = (1/s) I − Q diag(λ / (s (λ + s))) Qᵀ (Woodbury / projector–complement split)
  • lpinverse: pseudo-inverse using the same spectral formula with safe handling of zero/near-zero modes.
  • lsqrt: (s I + A)^{1/2} = Q diag(√(λ + s)) Qᵀ
  • lcholesky: returns a factor L with L Lᵀ = s I + A, namely L = Q diag(√(λ + s)) (orthonormal “spectral” factor)
  • llog: log(s I + A) = Q diag(log(λ + s)) Qᵀ
  • lpow: (s I + A)^p = Q diag((λ + s)^p) Qᵀ
  • diagonal: diag(s I + A) = s · 1 + diag(A)
  • ltrace: tr(s I + A) = s·n + tr(A) (with Hutchinson if needed)
  • lexp: exp(s I + A) = Q diag(exp(λ + s)) Qᵀ
Projector / anti-projector view

When leigh returns a truncated eigenspace Q ∈ ℝ^{n×k} (k ≤ n), let P := Q Qᵀ be the projector onto the retained subspace and P⊥ := I − P the orthogonal complement. Then

(s I + A)⁻¹
= Q (Λ + s I)⁻¹ Qᵀ  +  (1/s) P⊥,

i.e. the inverse acts as (Λ + s I)⁻¹ on span(Q) and as (1/s) I on its orthogonal complement. The implementation of linverse uses the equivalent Woodbury form

(s I + A)⁻¹ = (1/s) [ I − Q diag(λ / (λ + s)) Qᵀ ].

If leigh is full-rank, then P = I and P⊥ = 0, which recovers the usual full spectral formulas.

Caching notes
  • Q and S (eigenvectors/eigenvalues) are cached lazily by _ensure_eigh(). Any operation that changes the operator should call _invalidate_cache().
  • projector (Q Qᵀ) and complement (I − Q Qᵀ) are also cached on demand.
Arguments:

s : jax.Array Scalar added to the diagonal (isotropic shift). May be wrapped into a scalar ScaledLinearOperator(Identity, s). A : LinearOperator Symmetric linear operator (square).

Symmetry is **required and checked exactly**, including under
``jax.jit`` -- where it becomes a runtime error rather than a
trace-time one. The check costs nothing asymptotically: it runs in
:meth:`_ensure_eigh`, which is about to densify ``A`` for ``leigh``
regardless.

``_matmul`` and ``_todense`` stay permissive, since ``s*I + A`` is
computed correctly for any square ``A`` and ``smart_add`` routes every
``Identity + op`` sum through this class. Only the eigh-backed
shortcuts refuse.

Symmetry alone is enough for ``inverse``, ``eigh`` and ``slogdet``.
``sqrt``, ``lcholesky``, ``llog`` and fractional ``lpow`` additionally
require the *shifted* spectrum ``s + lambda`` to be non-negative
(strictly positive for the logarithm); that is checked per-operation
against the eigenvalues, which have already been computed by then.

Returns:

Type Description
A LinearOperator supporting matrix-free application and spectral transforms
of ``s I + A`` via the multipledispatch functions listed above.
-------
Example:

n = 100 s = jnp.array(0.1) A = utils.as_linop(jnp.diag(jnp.linspace(0.0, 5.0, n))) # symmetric L = IsotropicAdditiveLinearOperator(s, A) x = jnp.ones((n,)) y = (linverse(L) @ x) # apply (s I + A)^{-1} to a vector d = diagonal(L) # exact diagonal z = (lsqrt(L) @ x) # apply (s I + A)^{1/2} to a vector

Source code in linox/operators/isotropic.py 
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
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
class IsotropicAdditiveLinearOperator(AddLinearOperator):
    r"""Isotropic additive linear operator for matrices of the form.

        A_iso := s I + A,

    where ``s`` is a scalar (or a 0-arg scalar LinearOperator) and ``A`` is a
    symmetric LinearOperator. This class exposes fast, matrix-free implementations
    of common spectral transforms (inverse, pseudo-inverse, square root, log,
    powers, exp, Cholesky-like factor) by working in the eigenbasis of ``A``.

    ----------
    Core idea
    ----------
    If ``A = Q Λ Qᵀ`` is an eigendecomposition of ``A`` (with Λ diagonal and
    ``Qᵀ Q = I``), then

        s I + A = Q (Λ + s I) Qᵀ,

    so any spectral function ``f`` (e.g. inverse, sqrt, log, power, exp) satisfies

        f(s I + A) = Q f(Λ + s I) Qᵀ,

    which reduces the linear-algebra to elementwise operations on the eigenvalues.

    This class computes/caches an (optionally truncated) eigendecomposition via
    ``leigh(A)`` and then dispatches the following:

    * ``linverse``:       (s I + A)⁻¹ = (1/s) I − Q diag(λ / (s (λ + s))) Qᵀ
                          (Woodbury / projector–complement split)
    * ``lpinverse``:      pseudo-inverse using the same spectral formula with
                          safe handling of zero/near-zero modes.
    * ``lsqrt``:          (s I + A)^{1/2} = Q diag(√(λ + s)) Qᵀ
    * ``lcholesky``:      returns a factor L with L Lᵀ = s I + A, namely
                          L = Q diag(√(λ + s))   (orthonormal “spectral” factor)
    * ``llog``:           log(s I + A) = Q diag(log(λ + s)) Qᵀ
    * ``lpow``:           (s I + A)^p = Q diag((λ + s)^p) Qᵀ
    * ``diagonal``:       diag(s I + A) = s · 1 + diag(A)
    * ``ltrace``:         tr(s I + A) = s·n + tr(A)  (with Hutchinson if needed)
    * ``lexp``:           exp(s I + A) = Q diag(exp(λ + s)) Qᵀ

    -------------------------------
    Projector / anti-projector view
    -------------------------------
    When ``leigh`` returns a **truncated** eigenspace ``Q ∈ ℝ^{n×k}`` (k ≤ n),
    let P := Q Qᵀ be the projector onto the retained subspace and
    P⊥ := I − P the orthogonal complement. Then

        (s I + A)⁻¹
        = Q (Λ + s I)⁻¹ Qᵀ  +  (1/s) P⊥,

    i.e. the inverse acts as ``(Λ + s I)⁻¹`` on span(Q) and as ``(1/s) I`` on
    its orthogonal complement. The implementation of ``linverse`` uses the
    equivalent Woodbury form

        (s I + A)⁻¹ = (1/s) [ I − Q diag(λ / (λ + s)) Qᵀ ].

    If ``leigh`` is **full-rank**, then P = I and P⊥ = 0, which recovers the
    usual full spectral formulas.

    -------------
    Caching notes
    -------------
    * ``Q`` and ``S`` (eigenvectors/eigenvalues) are cached lazily by
      ``_ensure_eigh()``. Any operation that changes the operator should call
      ``_invalidate_cache()``.
    * ``projector`` (Q Qᵀ) and ``complement`` (I − Q Qᵀ) are also cached on demand.

    ----------

    Arguments:
    ----------
    s : jax.Array
        Scalar added to the diagonal (isotropic shift). May be wrapped into a
        scalar ``ScaledLinearOperator(Identity, s)``.
    A : LinearOperator
        Symmetric linear operator (square).

        Symmetry is **required and checked exactly**, including under
        ``jax.jit`` -- where it becomes a runtime error rather than a
        trace-time one. The check costs nothing asymptotically: it runs in
        :meth:`_ensure_eigh`, which is about to densify ``A`` for ``leigh``
        regardless.

        ``_matmul`` and ``_todense`` stay permissive, since ``s*I + A`` is
        computed correctly for any square ``A`` and ``smart_add`` routes every
        ``Identity + op`` sum through this class. Only the eigh-backed
        shortcuts refuse.

        Symmetry alone is enough for ``inverse``, ``eigh`` and ``slogdet``.
        ``sqrt``, ``lcholesky``, ``llog`` and fractional ``lpow`` additionally
        require the *shifted* spectrum ``s + lambda`` to be non-negative
        (strictly positive for the logarithm); that is checked per-operation
        against the eigenvalues, which have already been computed by then.

    -------

    Returns
    -------
    A LinearOperator supporting matrix-free application and spectral transforms
    of ``s I + A`` via the multipledispatch functions listed above.

    -------

    Example:
    -------
    >>> n = 100
    >>> s = jnp.array(0.1)
    >>> A = utils.as_linop(jnp.diag(jnp.linspace(0.0, 5.0, n)))  # symmetric
    >>> L = IsotropicAdditiveLinearOperator(s, A)
    >>> x = jnp.ones((n,))
    >>> y = (linverse(L) @ x)          # apply (s I + A)^{-1} to a vector
    >>> d = diagonal(L)                 # exact diagonal
    >>> z = (lsqrt(L) @ x)              # apply (s I + A)^{1/2} to a vector

    """

    def __init__(self, s: jax.Array, A: LinearOperator) -> None:
        self._A = utils.as_linop(A)
        if self._A.shape[-1] != self._A.shape[-2]:
            msg = "A must be a square matrix."
            raise ValueError(msg)
        self._s = ScaledLinearOperator(Identity(self._A.shape[0], dtype=self._A.dtype), s)
        self._Q = None
        self._S = None
        self._projector = None
        self._complement = None
        super().__init__(self._s, self._A)

    def _ensure_eigh(self) -> None:
        if (self._S is None) or (self._Q is None):
            # Guard here rather than in __init__: `_matmul`/`_todense` compute
            # s*I + A correctly for any square A, and `smart_add` rewrites
            # every `Identity + op` sum into this class. Only the eigh-based
            # shortcuts require symmetry, so only they need to refuse.
            _require_symmetric(self._A)
            self._S, self._Q = leigh(self._A)
            # invalidate derived caches
            self._projector = None
            self._complement = None

    def _invalidate_cache(self) -> None:
        self._Q = self._S = self._projector = self._complement = None

    @property
    def s(self) -> jax.Array:
        """Scalar operator component (s * I)."""
        return self._s

    @property
    def scalar(self) -> jax.Array:
        """Scalar value s from the isotropic shift."""
        return self._s.scalar

    @property
    def shape(self) -> tuple[int, int]:
        """Shape of the operator."""
        return self._A.shape

    @property
    def operator(self) -> LinearOperator:
        """The base linear operator A."""
        return self._A

    @property
    def Q(self) -> LinearOperator:
        """Eigenvectors of A (computed lazily via leigh)."""
        self._ensure_eigh()
        return self._Q

    @property
    def S(self) -> LinearOperator:
        """Eigenvalues of A (computed lazily via leigh)."""
        self._ensure_eigh()
        return self._S

    @property
    def projector(self) -> LinearOperator:
        """Projector onto the eigenspace Q Q^T (cached)."""
        self._ensure_eigh()
        if self._projector is None:
            self._projector = self._Q @ self._Q.T
        return self._projector

    @property
    def complement(self) -> LinearOperator:
        """Orthogonal complement projector I - Q Q^T (cached)."""
        self._ensure_eigh()
        if self._complement is None:
            self._complement = Identity(self.shape[0], dtype=self._A.dtype) - self.projector
        return self._complement

    def _matmul(self, arr: jax.Array):
        return self._s @ arr + self._A @ arr

    def _todense(self) -> jax.Array:
        return self._s._todense() + self._A._todense()

    def tree_flatten(self) -> tuple[tuple, dict]:
        """Flatten this operator into JAX pytree children and static data."""
        children = (self._s.scalar, self._A)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        """Reconstruct this operator from JAX pytree children and static data."""
        s, A = children
        return cls(s, A)

Q: LinearOperator property

Eigenvectors of A (computed lazily via leigh).

S: LinearOperator property

Eigenvalues of A (computed lazily via leigh).

complement: LinearOperator property

Orthogonal complement projector I - Q Q^T (cached).

operator: LinearOperator property

The base linear operator A.

projector: LinearOperator property

Projector onto the eigenspace Q Q^T (cached).

s: jax.Array property

Scalar operator component (s * I).

scalar: jax.Array property

Scalar value s from the isotropic shift.

shape: tuple[int, int] property

Shape of the operator.

tree_flatten() -> tuple[tuple, dict]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/isotropic.py 
327
328
329
330
331
def tree_flatten(self) -> tuple[tuple, dict]:
    """Flatten this operator into JAX pytree children and static data."""
    children = (self._s.scalar, self._A)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data, children) classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/isotropic.py 
333
334
335
336
337
@classmethod
def tree_unflatten(cls, aux_data, children):
    """Reconstruct this operator from JAX pytree children and static data."""
    s, A = children
    return cls(s, A)

Low rank representations as linear operators.

This module implements various low rank representations as linear operators, including:

  • :class:LowRank: Represents a low rank matrix :math:A = U \text{diag}(S) V^T
  • :class:SymmetricLowRank: Represents a symmetric low rank matrix :math:A = U \text{diag}(S) U^T
  • :class:IsotropicScalingPlusSymmetricLowRank: Represents :math:\sigma I + U \text{diag}(S) U^T
  • :class:PositiveDiagonalPlusSymmetricLowRank: Represents :math:D + \alpha U \text{diag}(S) U^T where :math:D is a positive diagonal matrix

IsotropicScalingPlusSymmetricLowRank

Bases: AddLinearOperator

Isotropic scaling plus symmetric low rank operator.

For scalar :math:\sigma, matrix :math:U, and vector :math:S, this represents :math:A = \sigma I + U \text{diag}(S) U^T. The action on a vector :math:x is given by :math:Ax = \sigma x + U(S \odot (U^T x)) where :math:\odot denotes element-wise multiplication.

Args: scalar: Isotropic scaling factor :math:\sigma U: Factor matrix S: Vector of singular values

Source code in linox/operators/lowrank.py 
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
class IsotropicScalingPlusSymmetricLowRank(AddLinearOperator):
    r"""Isotropic scaling plus symmetric low rank operator.

    For scalar :math:`\sigma`, matrix :math:`U`, and vector :math:`S`, this represents
    :math:`A = \sigma I + U \text{diag}(S) U^T`. The action on a vector :math:`x` is
    given by :math:`Ax = \sigma x + U(S \odot (U^T x))` where :math:`\odot` denotes
    element-wise multiplication.

    Args:
        scalar: Isotropic scaling factor :math:`\sigma`
        U: Factor matrix
        S: Vector of singular values
    """

    def __init__(self, scalar: jax.Array, U: jax.Array, S: jax.Array) -> None:
        self._scalar = scalar

        self._U = U
        self._S = S

        # Move to an abstract base class instead
        super().__init__(
            self._scalar * Identity(self._U.shape[-2], dtype=self._U.dtype),
            SymmetricLowRank(self._U, self._S),
        )

        # Add general tagging
        # self.is_symmetric = True
        self._lr_eigh = (self._U, self._S)

    @property
    def scalar(self) -> float:
        """Isotropic scaling factor."""
        return self._scalar

    @property
    def U(self) -> jax.Array:
        """Factor matrix."""
        return self._U

    @property
    def S(self) -> jax.Array:
        """Singular values vector."""
        return self._S

    @property
    def lr_eigh(self) -> jax.Array:
        """Low rank eigendecomposition (U, S)."""
        return self._lr_eigh

    def transpose(self) -> "IsotropicScalingPlusSymmetricLowRank":
        """Return transpose (self for symmetric operators)."""
        return self

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration."""
        # We need to override the AddLinearOperator's tree_flatten
        children = (self._scalar, self._U, self._S)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "IsotropicScalingPlusSymmetricLowRank":
        """Unflatten for JAX pytree registration."""
        del aux_data
        scalar, U, S = children
        return cls(scalar=scalar, U=U, S=S)

S: jax.Array property

Singular values vector.

U: jax.Array property

Factor matrix.

lr_eigh: jax.Array property

Low rank eigendecomposition (U, S).

scalar: float property

Isotropic scaling factor.

transpose() -> IsotropicScalingPlusSymmetricLowRank

Return transpose (self for symmetric operators).

Source code in linox/operators/lowrank.py 
195
196
197
def transpose(self) -> "IsotropicScalingPlusSymmetricLowRank":
    """Return transpose (self for symmetric operators)."""
    return self

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

Source code in linox/operators/lowrank.py 
199
200
201
202
203
204
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration."""
    # We need to override the AddLinearOperator's tree_flatten
    children = (self._scalar, self._U, self._S)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> IsotropicScalingPlusSymmetricLowRank classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/lowrank.py 
206
207
208
209
210
211
212
213
214
215
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "IsotropicScalingPlusSymmetricLowRank":
    """Unflatten for JAX pytree registration."""
    del aux_data
    scalar, U, S = children
    return cls(scalar=scalar, U=U, S=S)

LowRank

Bases: LinearOperator

Low rank operator.

For matrices :math:U, :math:S, and :math:V, this represents the low rank matrix :math:A = U \text{diag}(S) V^T. The action on a vector :math:x is given by :math:Ax = U(S \odot (V^T x)) where :math:\odot denotes element-wise multiplication.

Args: U: Left factor matrix S: Vector of singular values (optional, defaults to ones) V: Right factor matrix (optional, defaults to U)

Source code in linox/operators/lowrank.py 
 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
class LowRank(LinearOperator):
    r"""Low rank operator.

    For matrices :math:`U`, :math:`S`, and :math:`V`, this represents the low rank
    matrix :math:`A = U \text{diag}(S) V^T`. The action on a vector :math:`x` is given
    by :math:`Ax = U(S \odot (V^T x))` where :math:`\odot` denotes element-wise
    multiplication.

    Args:
        U: Left factor matrix
        S: Vector of singular values (optional, defaults to ones)
        V: Right factor matrix (optional, defaults to U)
    """

    def __init__(self, U: jax.Array, S: jax.Array | None = None, V: jax.Array | None = None) -> None:
        # Check shapes
        if S is not None:
            assert U.shape[-1] == S.shape[-1]
        if V is not None:
            assert U.shape[-1] == V.shape[-1]

        self._U = U
        self._S = S if S is not None else jnp.ones(U.shape[-1])
        self._V = V

        # Move to an abstract base class instead
        super().__init__(shape=(*U.shape[:-2], U.shape[-2], U.shape[-2]), dtype=U.dtype)

    @property
    def U(self) -> jax.Array:
        """Left factor matrix."""
        return self._U

    @property
    def S(self) -> jax.Array:
        """Singular values vector."""
        return self._S

    @property
    def V(self) -> jax.Array:
        """Right factor matrix (defaults to U if not provided)."""
        return self._V if self._V is not None else self._U

    def _matmul(self, arr: jnp.array) -> jnp.array:
        return self.U @ (self.S[:, None] * (self.V.T @ arr))

    def _todense(self) -> jnp.array:
        return self.U @ jnp.diag(self.S) @ self.V.T

    def transpose(self) -> "LowRank":
        """Return transposed low rank operator."""
        return LowRank(self.V, self.S, self.U)

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration."""
        children = (self._U, self._S, self._V)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "LowRank":
        """Unflatten for JAX pytree registration."""
        del aux_data
        U, S, V = children
        return cls(U=U, S=S, V=V)

S: jax.Array property

Singular values vector.

U: jax.Array property

Left factor matrix.

V: jax.Array property

Right factor matrix (defaults to U if not provided).

transpose() -> LowRank

Return transposed low rank operator.

Source code in linox/operators/lowrank.py 
85
86
87
def transpose(self) -> "LowRank":
    """Return transposed low rank operator."""
    return LowRank(self.V, self.S, self.U)

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

Source code in linox/operators/lowrank.py 
89
90
91
92
93
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration."""
    children = (self._U, self._S, self._V)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> LowRank classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/lowrank.py 
 95
 96
 97
 98
 99
100
101
102
103
104
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "LowRank":
    """Unflatten for JAX pytree registration."""
    del aux_data
    U, S, V = children
    return cls(U=U, S=S, V=V)

PositiveDiagonalPlusSymmetricLowRank

Bases: AddLinearOperator

Positive diagonal plus symmetric low rank operator.

For positive diagonal matrix :math:D, matrix :math:U, vector :math:S, and scalar :math:\alpha, this represents :math:A = D + \alpha U \text{diag}(S) U^T. The action on a vector :math:x is given by :math:Ax = Dx + \alpha U(S \odot (U^T x)) where :math:\odot denotes element-wise multiplication.

Args: diagonal: Positive diagonal matrix :math:D low_rank: Symmetric low rank component :math:U \text{diag}(S) U^T low_rank_scale: Scaling factor :math:\alpha (default: 1.0)

Source code in linox/operators/lowrank.py 
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
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
class PositiveDiagonalPlusSymmetricLowRank(AddLinearOperator):
    r"""Positive diagonal plus symmetric low rank operator.

    For positive diagonal matrix :math:`D`, matrix :math:`U`, vector :math:`S`, and
    scalar :math:`\alpha`, this represents :math:`A = D + \alpha U \text{diag}(S) U^T`.
    The action on a vector :math:`x` is given by
    :math:`Ax = Dx + \alpha U(S \odot (U^T x))` where :math:`\odot` denotes
    element-wise multiplication.

    Args:
        diagonal: Positive diagonal matrix :math:`D`
        low_rank: Symmetric low rank component :math:`U \text{diag}(S) U^T`
        low_rank_scale: Scaling factor :math:`\alpha` (default: 1.0)
    """

    def __init__(
        self,
        diagonal: Diagonal,  # D
        low_rank: SymmetricLowRank,  # U S U^T
        low_rank_scale: float = 1.0,  # a
    ) -> None:
        self._diagonal = diagonal
        self._low_rank = low_rank
        self._low_rank_scale = low_rank_scale

        super().__init__(self._diagonal, self._low_rank_scale * self._low_rank)

    @property
    def diagonal(self) -> jax.Array:
        """Diagonal component."""
        return self._diagonal

    @property
    def low_rank(self) -> SymmetricLowRank:
        """Low rank component."""
        return self._low_rank

    @property
    def low_rank_scale(self) -> float:
        """Scaling factor for low rank component."""
        return self._low_rank_scale

    @functools.cached_property
    def _id_plus_low_rank(self) -> IsotropicScalingPlusSymmetricLowRank:
        """1 + a (D^{-1/2} U) S (D^{-1/2} U)^T = D^{-1/2} (D + a U S U^T) D^{-1/2}."""
        U, sqrt_S, _ = jnp.linalg.svd(
            ((self.low_rank.U * jnp.sqrt(self.low_rank.S)) / jnp.sqrt(self._diagonal.diag[:, None])),
            full_matrices=False,
            compute_uv=True,
        )

        return IsotropicScalingPlusSymmetricLowRank(
            1.0,
            U,
            self._low_rank_scale * sqrt_S**2,
        )

    def transpose(self) -> "PositiveDiagonalPlusSymmetricLowRank":
        """Return transpose (self for symmetric operators)."""
        return self

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration."""
        # We need to override the AddLinearOperator's tree_flatten
        children = (self._diagonal, self._low_rank, self._low_rank_scale)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "PositiveDiagonalPlusSymmetricLowRank":
        """Unflatten for JAX pytree registration."""
        del aux_data
        diagonal, low_rank, low_rank_scale = children
        return cls(diagonal=diagonal, low_rank=low_rank, low_rank_scale=low_rank_scale)

diagonal: jax.Array property

Diagonal component.

low_rank: SymmetricLowRank property

Low rank component.

low_rank_scale: float property

Scaling factor for low rank component.

transpose() -> PositiveDiagonalPlusSymmetricLowRank

Return transpose (self for symmetric operators).

Source code in linox/operators/lowrank.py 
311
312
313
def transpose(self) -> "PositiveDiagonalPlusSymmetricLowRank":
    """Return transpose (self for symmetric operators)."""
    return self

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

Source code in linox/operators/lowrank.py 
315
316
317
318
319
320
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration."""
    # We need to override the AddLinearOperator's tree_flatten
    children = (self._diagonal, self._low_rank, self._low_rank_scale)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> PositiveDiagonalPlusSymmetricLowRank classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/lowrank.py 
322
323
324
325
326
327
328
329
330
331
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "PositiveDiagonalPlusSymmetricLowRank":
    """Unflatten for JAX pytree registration."""
    del aux_data
    diagonal, low_rank, low_rank_scale = children
    return cls(diagonal=diagonal, low_rank=low_rank, low_rank_scale=low_rank_scale)

SymmetricLowRank

Bases: LowRank

Symmetric low rank operator.

For matrices :math:U and :math:S, this represents the symmetric low rank matrix :math:A = U \text{diag}(S) U^T. The action on a vector :math:x is given by :math:Ax = U(S \odot (U^T x)) where :math:\odot denotes element-wise multiplication.

Args: U: Factor matrix S: Vector of singular values (optional, defaults to ones)

Source code in linox/operators/lowrank.py 
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
class SymmetricLowRank(LowRank):
    r"""Symmetric low rank operator.

    For matrices :math:`U` and :math:`S`, this represents the symmetric low rank matrix
    :math:`A = U \text{diag}(S) U^T`. The action on a vector :math:`x` is given by
    :math:`Ax = U(S \odot (U^T x))` where :math:`\odot` denotes element-wise
    multiplication.

    Args:
        U: Factor matrix
        S: Vector of singular values (optional, defaults to ones)
    """

    def __init__(self, U: jax.Array, S: jax.Array | None = None) -> None:
        super().__init__(U, S)

    def transpose(self) -> "SymmetricLowRank":
        """Return transpose (self for symmetric operators)."""
        return self

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration."""
        children = (self._U, self._S)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "SymmetricLowRank":
        """Unflatten for JAX pytree registration."""
        del aux_data
        U, S = children
        return cls(U=U, S=S)

transpose() -> SymmetricLowRank

Return transpose (self for symmetric operators).

Source code in linox/operators/lowrank.py 
123
124
125
def transpose(self) -> "SymmetricLowRank":
    """Return transpose (self for symmetric operators)."""
    return self

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

Source code in linox/operators/lowrank.py 
127
128
129
130
131
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration."""
    children = (self._U, self._S)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> SymmetricLowRank classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/lowrank.py 
133
134
135
136
137
138
139
140
141
142
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "SymmetricLowRank":
    """Unflatten for JAX pytree registration."""
    del aux_data
    U, S = children
    return cls(U=U, S=S)

_(A: PositiveDiagonalPlusSymmetricLowRank) -> tuple[jax.Array, jax.Array]

Compute log-determinant using the Matrix Determinant Lemma.

|D + U S U^T| = |D| * |I + S U^T D^{-1} U| or |I + U S U^T D^{-1}| = |I + S U^T D^{-1} U| (Sylvester's determinant theorem)

Actually: |D + U S U^T| = |D| |I + S^{-1} + U^T D^{-1} U| wait no. Matrix Determinant Lemma: |A + U W V^T| = |D| |W| |W^{-1} + V^T A^{-1} U| Here A=D, W=S (diag), V=U, U=U. |D + U S U^T| = |D| |S| |S^{-1} + U^T D^{-1} U|

But S might be singular or non-invertible if used loosely. Alternative: |D + U S U^T| = |D| |I + S U^T D^{-1} U| Currently we assume S is diagonal vector.

C = S * (U^T D^{-1} U) inner = I + C

Source code in linox/operators/lowrank.py 
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
@slogdet.dispatch
def _(A: PositiveDiagonalPlusSymmetricLowRank) -> tuple[jax.Array, jax.Array]:
    r"""Compute log-determinant using the Matrix Determinant Lemma.

    |D + U S U^T| = |D| * |I + S U^T D^{-1} U|
    or |I + U S U^T D^{-1}| = |I + S U^T D^{-1} U| (Sylvester's determinant theorem)

    Actually: |D + U S U^T| = |D| |I + S^{-1} + U^T D^{-1} U| wait no.
    Matrix Determinant Lemma: |A + U W V^T| = |D| |W| |W^{-1} + V^T A^{-1} U|
    Here A=D, W=S (diag), V=U, U=U.
    |D + U S U^T| = |D| |S| |S^{-1} + U^T D^{-1} U|

    But S might be singular or non-invertible if used loosely.
    Alternative: |D + U S U^T| = |D| |I + S U^T D^{-1} U|
    Currently we assume S is diagonal vector.

    C = S * (U^T D^{-1} U)
    inner = I + C
    """
    from linox.operators.arithmetic import slogdet

    sign_D, logdet_D = slogdet(A.diagonal)

    # Inner matrix: I + alpha * S * (U^T D^{-1} U)
    # Shape k x k
    D_inv = linverse(A.diagonal)

    # U^T D^{-1} U
    # If D is diagonal, D_inv is diagonal.
    # We can compute efficiently: (U / D_diag)^T @ U
    Ut_Dinv_U = (A.low_rank.U.T @ D_inv) @ A.low_rank.U

    # Scale by S * alpha
    # S is vector, acting as diagonal matrix
    scaled = A.low_rank.S[:, None] * Ut_Dinv_U  # S @ (Ut Dinv U)

    # Add Identity
    k = A.low_rank.S.shape[0]
    inner = jnp.eye(k, dtype=A.dtype) + A.low_rank_scale * scaled

    sign_inner, logdet_inner = jnp.linalg.slogdet(inner)

    return sign_D * sign_inner, logdet_D + logdet_inner

Linear operator for representing an eigenvalue decomposition.

This module provides:

  • :class:EigenD: Represents a symmetric linear operator in its eigenvalue decomposition form :math:A = Q \Lambda Q^T where :math:Q is orthogonal and :math:\Lambda is diagonal. Both Q and Lambda are LinearOperators.

EigenD

Bases: LinearOperator

A linear operator representing an eigenvalue decomposition.

Represents :math:A = Q \Lambda Q^T where :math:Q is orthogonal and :math:\Lambda is diagonal. Both Q and Lambda are stored as LinearOperators, enabling structured representations (e.g., Kronecker products of eigenvectors).

Args: Q: Orthogonal matrix of eigenvectors (LinearOperator or array) Lambda: Diagonal matrix of eigenvalues (LinearOperator or array)

Example: >>> import jax.numpy as jnp >>> from linox import Matrix, Diagonal, leigh >>> A = Matrix(jnp.diag(jnp.array([1.0, 2.0, 3.0]))) >>> lam, Q = leigh(A) >>> eigend = EigenD(Q, Diagonal(lam)) >>> # eigend represents A = Q @ diag(lam) @ Q.T

Source code in linox/operators/eigen.py 
 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
class EigenD(LinearOperator):
    r"""A linear operator representing an eigenvalue decomposition.

    Represents :math:`A = Q \Lambda Q^T` where :math:`Q` is orthogonal and
    :math:`\Lambda` is diagonal. Both Q and Lambda are stored as LinearOperators,
    enabling structured representations (e.g., Kronecker products of eigenvectors).

    Args:
        Q: Orthogonal matrix of eigenvectors (LinearOperator or array)
        Lambda: Diagonal matrix of eigenvalues (LinearOperator or array)

    Example:
        >>> import jax.numpy as jnp
        >>> from linox import Matrix, Diagonal, leigh
        >>> A = Matrix(jnp.diag(jnp.array([1.0, 2.0, 3.0])))
        >>> lam, Q = leigh(A)
        >>> eigend = EigenD(Q, Diagonal(lam))
        >>> # eigend represents A = Q @ diag(lam) @ Q.T
    """

    def __init__(
        self,
        Q: LinearOperator | jax.Array,
        Lambda: LinearOperator | jax.Array,
    ) -> None:
        self._Q = as_linop(Q)
        if isinstance(Lambda, jax.Array) and Lambda.ndim == 1:
            self._Lambda = Diagonal(Lambda)
        else:
            self._Lambda = as_linop(Lambda)

        n = self._Q.shape[0]
        super().__init__(shape=(n, n), dtype=self._Q.dtype)

    @property
    def Q(self) -> LinearOperator:
        """The orthonormal eigenvector operator ``Q``."""
        return self._Q

    @property
    def Lambda(self) -> LinearOperator:
        """The eigenvalue operator ``Lambda``."""
        return self._Lambda

    @property
    def eigenvalues(self) -> jax.Array:
        """The eigenvalues, as a :class:`jax.Array`."""
        return diagonal(self._Lambda)

    def _matmul(self, vec: jax.Array) -> jax.Array:
        return self._Q @ (self._Lambda @ (self._Q.T @ vec))

    def _todense(self) -> jax.Array:
        Q_dense = self._Q._todense()
        lam = self.eigenvalues
        return Q_dense @ (lam[:, None] * Q_dense.T)

    def transpose(self) -> "EigenD":
        """Return the transpose of this operator."""
        return self

    def tree_flatten(self) -> tuple[tuple, dict]:
        """Flatten this operator into JAX pytree children and static data."""
        children = (self._Q, self._Lambda)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict,
        children: tuple,
    ) -> "EigenD":
        """Reconstruct this operator from JAX pytree children and static data."""
        Q, Lambda = children
        return cls(Q=Q, Lambda=Lambda)

Lambda: LinearOperator property

The eigenvalue operator Lambda.

Q: LinearOperator property

The orthonormal eigenvector operator Q.

eigenvalues: jax.Array property

The eigenvalues, as a :class:jax.Array.

transpose() -> EigenD

Return the transpose of this operator.

Source code in linox/operators/eigen.py 
85
86
87
def transpose(self) -> "EigenD":
    """Return the transpose of this operator."""
    return self

tree_flatten() -> tuple[tuple, dict]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/eigen.py 
89
90
91
92
93
def tree_flatten(self) -> tuple[tuple, dict]:
    """Flatten this operator into JAX pytree children and static data."""
    children = (self._Q, self._Lambda)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict, children: tuple) -> EigenD classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/eigen.py 
 95
 96
 97
 98
 99
100
101
102
103
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict,
    children: tuple,
) -> "EigenD":
    """Reconstruct this operator from JAX pytree children and static data."""
    Q, Lambda = children
    return cls(Q=Q, Lambda=Lambda)

Block matrix operations for linear operators.

This module implements various block matrix operations for linear operators, including:

  • :class:BlockMatrix: Represents a general block matrix :math:\begin{bmatrix} A_{11} & \cdots & A_{1n} \\ \vdots & \ddots & \vdots \\ A_{m1} & \cdots & A_{mn} \end{bmatrix}
  • :class:BlockMatrix2x2: Represents a 2x2 block matrix :math:\begin{bmatrix} A & B \\ C & D \end{bmatrix}
  • :class:BlockDiagonal: Represents a block diagonal matrix :math:\begin{bmatrix} A_1 & & \\ & \ddots & \\ & & A_n \end{bmatrix}

These operators allow efficient representation and computation of structured linear transformations through block matrix operations.

BlockDiagonal

Bases: LinearOperator

Block diagonal matrix.

For linear operators :math:A_1, \ldots, A_n, this represents the block diagonal matrix :math:\begin{bmatrix} A_1 & & \\ & \ddots & \\ & & A_n \end{bmatrix} where :math:\begin{bmatrix} A_1 & & \\ & \ddots & \\ & & A_n \end{bmatrix} \begin{bmatrix} x_1 \\ \vdots \\ x_n \end{bmatrix} = \begin{bmatrix} A_1x_1 \\ \vdots \\ A_nx_n \end{bmatrix}

Args: *blocks: Variable number of linear operators to be placed on the diagonal. All blocks must have compatible shapes.

Source code in linox/operators/block.py 
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
class BlockDiagonal(LinearOperator):
    r"""Block diagonal matrix.

    For linear operators :math:`A_1, \ldots, A_n`, this represents the block diagonal
    matrix :math:`\begin{bmatrix} A_1 & & \\ & \ddots & \\ & & A_n \end{bmatrix}`
    where :math:`\begin{bmatrix} A_1 & & \\ & \ddots & \\ & & A_n \end{bmatrix}
    \begin{bmatrix} x_1 \\ \vdots \\ x_n \end{bmatrix} =
    \begin{bmatrix} A_1x_1 \\ \vdots \\ A_nx_n \end{bmatrix}`

    Args:
        *blocks: Variable number of linear operators to be placed on the diagonal.
            All blocks must have compatible shapes.
    """

    def __init__(self, *blocks: LinearOperator) -> None:
        if len(blocks) < 1:
            msg = "At least one block must be given."
            raise ValueError(msg)

        self.blocks = [utils.as_linop(block) for block in blocks]
        self._all_blocks_square = all(block.shape[0] == block.shape[1] for block in blocks)

        dtype = reduce(jnp.promote_types, (block.dtype for block in blocks))
        shape_0 = sum(block.shape[0] for block in blocks)
        shape_1 = sum(block.shape[1] for block in blocks)
        # Plain Python ints, not a jnp array: `jnp.split` needs concrete
        # indices, and a traced array here made every BlockDiagonal matvec
        # fail under `jax.jit`. The block shapes are static anyway.
        self.split_indices = tuple(itertools.accumulate(block.shape[1] for block in blocks))[:-1]

        super().__init__((shape_0, shape_1), dtype)

    def _split_input(self, x: jax.Array) -> list[jax.Array]:
        return jnp.split(x, self.split_indices, axis=-2)

    def _matmul(self, x: jax.Array) -> jax.Array:
        res = jnp.concatenate(
            [block @ cur_x for block, cur_x in zip(self.blocks, self._split_input(x), strict=False)],
            axis=-2,
        )
        return res

    def _todense(self) -> jax.Array:
        return jax.scipy.linalg.block_diag(*[block._todense() for block in self.blocks])

    def transpose(self) -> "BlockDiagonal":
        """Return the transpose of this operator."""
        return BlockDiagonal(*[block.transpose() for block in self.blocks])

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = tuple(self.blocks)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "BlockDiagonal":
        """Reconstruct this operator from JAX pytree children and static data."""
        del aux_data
        return cls(*children)

transpose() -> BlockDiagonal

Return the transpose of this operator.

Source code in linox/operators/block.py 
280
281
282
def transpose(self) -> "BlockDiagonal":
    """Return the transpose of this operator."""
    return BlockDiagonal(*[block.transpose() for block in self.blocks])

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/block.py 
284
285
286
287
288
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = tuple(self.blocks)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> BlockDiagonal classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/block.py 
290
291
292
293
294
295
296
297
298
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "BlockDiagonal":
    """Reconstruct this operator from JAX pytree children and static data."""
    del aux_data
    return cls(*children)

BlockMatrix

Bases: LinearOperator

A block matrix where each block is represented by a linear operator.

For linear operators :math:A_{ij}, this represents the block matrix :math:\begin{bmatrix} A_{11} & \cdots & A_{1n} \\ \vdots & \ddots & \vdots \\ A_{m1} & \cdots & A_{mn} \end{bmatrix} where each block :math:A_{ij} is a linear operator of compatible shape.

Args: blocks: A nested list of LinearOperatorLike instances representing the blocks of the matrix. Shapes must be valid such that creating a block matrix is possible.

Source code in linox/operators/block.py 
 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
class BlockMatrix(LinearOperator):
    r"""A block matrix where each block is represented by a linear operator.

    For linear operators :math:`A_{ij}`, this represents the block matrix
    :math:`\begin{bmatrix} A_{11} & \cdots & A_{1n} \\
    \vdots & \ddots & \vdots \\ A_{m1} & \cdots & A_{mn} \end{bmatrix}`
    where each block :math:`A_{ij}` is a linear operator of compatible shape.

    Args:
        blocks: A nested list of `LinearOperatorLike` instances representing the blocks
            of the matrix. Shapes must be valid such that creating a block matrix
            is possible.
    """

    def __init__(self, blocks: list[list[LinearOperatorLike]]) -> None:
        self._blocks = [[utils.as_linop(x) for x in sub_list] for sub_list in blocks]
        self._block_shape = (len(self._blocks), len(self._blocks[0]))

        # Determine the dtype
        dtype = self._blocks[0][0].dtype
        for i, j in product(range(len(self._blocks)), range(len(self._blocks[0]))):
            assert self._blocks[i][j].dtype == dtype, "All blocks must have the same dtype."
            expected_shape = (self._blocks[i][0].shape[0], self._blocks[0][j].shape[1])
            if self._blocks[i][j].shape != expected_shape:
                msg = f"Shape error in block [{i}, {j}]: Expected shape {expected_shape}, got shape {self._blocks[i, j].shape}."
                raise ValueError(msg)

            # Compute the total shape of the block matrix
            num_rows = sum(row[0].shape[0] for row in self._blocks)
            num_cols = sum(block.shape[1] for block in self._blocks[0][:])

            # Initialize the super class
            super().__init__(shape=(num_rows, num_cols), dtype=dtype)

            # Store the column sizes instead of precomputing indices
            self._col_sizes = [block.shape[1] for block in self._blocks[0]]

    @property
    def blocks(self) -> jnp.ndarray:
        """The blocks of the block matrix."""
        return self._blocks

    def _split_input(self, x: jnp.ndarray) -> list[jnp.ndarray]:
        """Split the input into blocks using column sizes."""
        results = []
        start_idx = 0
        for size in self._col_sizes:
            results.append(x[..., start_idx : start_idx + size, :])
            start_idx += size
        return results

    def _matmul(self, arr: jnp.ndarray) -> jnp.ndarray:
        # Split the input according to the block structure
        arr_split = self._split_input(arr)
        row_wise_results = []

        # Perform matrix multiplication for each row of blocks
        for i in range(self._block_shape[0]):
            row_wise_results.append(
                jnp.sum(
                    jnp.array([block @ cur_x for block, cur_x in zip(self.blocks[i], arr_split, strict=False)]),
                    axis=0,
                )
            )

        # Concatenate the results to form the final matrix
        return jnp.concatenate(row_wise_results, axis=-2)

    def _todense(self) -> jax.Array:
        """Convert the block matrix to a dense matrix."""
        blocks = [[None for _ in range(self._block_shape[1])] for _ in range(self._block_shape[0])]
        for i, j in np.ndindex(self._block_shape):
            blocks[i][j] = self._blocks[i][j]._todense()
        return jnp.block(blocks)

    def transpose(self) -> "BlockMatrix":
        """Transpose the block matrix."""
        blocks_t = [[self._blocks[i][j].transpose() for i in range(self._block_shape[0])] for j in range(self._block_shape[1])]

        return BlockMatrix(blocks_t)

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        # Flatten the nested list of blocks into a single tuple
        """Flatten this operator into JAX pytree children and static data."""
        flattened_blocks = []
        for row in self._blocks:
            flattened_blocks.extend(row)

        children = tuple(flattened_blocks)
        aux_data = {"block_shape": self._block_shape, "col_sizes": self._col_sizes}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "BlockMatrix":
        """Reconstruct this operator from JAX pytree children and static data."""
        block_shape = aux_data["block_shape"]

        # Reconstruct the nested list structure
        blocks = []
        idx = 0
        for _i in range(block_shape[0]):
            row = []
            for _j in range(block_shape[1]):
                row.append(children[idx])
                idx += 1
            blocks.append(row)

        # Create the instance
        instance = cls(blocks=blocks)

        # If col_sizes was saved, restore it to avoid recomputation
        if "col_sizes" in aux_data:
            instance._col_sizes = aux_data["col_sizes"]

        return instance

blocks: jnp.ndarray property

The blocks of the block matrix.

transpose() -> BlockMatrix

Transpose the block matrix.

Source code in linox/operators/block.py 
108
109
110
111
112
def transpose(self) -> "BlockMatrix":
    """Transpose the block matrix."""
    blocks_t = [[self._blocks[i][j].transpose() for i in range(self._block_shape[0])] for j in range(self._block_shape[1])]

    return BlockMatrix(blocks_t)

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/block.py 
114
115
116
117
118
119
120
121
122
123
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    # Flatten the nested list of blocks into a single tuple
    """Flatten this operator into JAX pytree children and static data."""
    flattened_blocks = []
    for row in self._blocks:
        flattened_blocks.extend(row)

    children = tuple(flattened_blocks)
    aux_data = {"block_shape": self._block_shape, "col_sizes": self._col_sizes}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> BlockMatrix classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/block.py 
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
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "BlockMatrix":
    """Reconstruct this operator from JAX pytree children and static data."""
    block_shape = aux_data["block_shape"]

    # Reconstruct the nested list structure
    blocks = []
    idx = 0
    for _i in range(block_shape[0]):
        row = []
        for _j in range(block_shape[1]):
            row.append(children[idx])
            idx += 1
        blocks.append(row)

    # Create the instance
    instance = cls(blocks=blocks)

    # If col_sizes was saved, restore it to avoid recomputation
    if "col_sizes" in aux_data:
        instance._col_sizes = aux_data["col_sizes"]

    return instance

BlockMatrix2x2

Bases: LinearOperator

2x2 Block Matrix.

For linear operators :math:A, B, C, D, this represents the block matrix :math:\begin{bmatrix} A & B \\ C & D \end{bmatrix} where :math:\begin{bmatrix} A & B \\ C & D \end{bmatrix} \begin{bmatrix} x \\ y \end{bmatrix} = \begin{bmatrix} Ax + By \\ Cx + Dy \end{bmatrix}

Args: A: The top-left block of the matrix. B: The top-right block of the matrix. C: The bottom-left block of the matrix. D: The bottom-right block of the matrix.

Source code in linox/operators/block.py 
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
class BlockMatrix2x2(LinearOperator):
    r"""2x2 Block Matrix.

    For linear operators :math:`A, B, C, D`, this represents the block matrix
    :math:`\begin{bmatrix} A & B \\ C & D \end{bmatrix}`
    where :math:`\begin{bmatrix} A & B \\ C & D \end{bmatrix}
    \begin{bmatrix} x \\ y \end{bmatrix} =
    \begin{bmatrix} Ax + By \\ Cx + Dy \end{bmatrix}`

    Args:
        A: The top-left block of the matrix.
        B: The top-right block of the matrix.
        C: The bottom-left block of the matrix.
        D: The bottom-right block of the matrix.
    """

    def __init__(
        self,
        A: LinearOperatorLike,
        B: LinearOperatorLike,
        C: LinearOperatorLike,
        D: LinearOperatorLike,
    ) -> None:
        self.A = utils.as_linop(A)
        self.B = utils.as_linop(B)
        self.C = utils.as_linop(C)
        self.D = utils.as_linop(D)

        dtype = reduce(jnp.promote_types, (self.A.dtype, self.B.dtype, self.C.dtype, self.D.dtype))

        super().__init__(shape=(A.shape[0] + D.shape[0], A.shape[1] + D.shape[1]), dtype=dtype)

    def _split_input(self, arr: jnp.ndarray) -> jnp.ndarray:
        return jnp.split(arr, [self.A.shape[1]], axis=-2)

    def _matmul(self, arr: jnp.ndarray) -> jnp.ndarray:
        arr0, arr1 = self._split_input(arr)
        return jnp.concatenate(
            [
                self.A @ arr0 + self.B @ arr1,
                self.C @ arr0 + self.D @ arr1,
            ],
            axis=-2,
        )

    def _todense(self) -> jax.Array:
        """Convert the block matrix to a dense matrix."""
        A = self.A._todense()
        B = self.B._todense()
        C = self.C._todense()
        D = self.D._todense()

        return jnp.block([[A, B], [C, D]])

    def transpose(self) -> "BlockMatrix2x2":
        """Return the transpose of this operator."""
        return BlockMatrix2x2(
            A=self.A.transpose(),
            B=self.C.transpose(),
            C=self.B.transpose(),
            D=self.D.transpose(),
        )

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten this operator into JAX pytree children and static data."""
        children = (self.A, self.B, self.C, self.D)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "BlockMatrix2x2":
        """Reconstruct this operator from JAX pytree children and static data."""
        del aux_data
        A, B, C, D = children
        return cls(A=A, B=B, C=C, D=D)

transpose() -> BlockMatrix2x2

Return the transpose of this operator.

Source code in linox/operators/block.py 
208
209
210
211
212
213
214
215
def transpose(self) -> "BlockMatrix2x2":
    """Return the transpose of this operator."""
    return BlockMatrix2x2(
        A=self.A.transpose(),
        B=self.C.transpose(),
        C=self.B.transpose(),
        D=self.D.transpose(),
    )

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/block.py 
217
218
219
220
221
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten this operator into JAX pytree children and static data."""
    children = (self.A, self.B, self.C, self.D)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> BlockMatrix2x2 classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/block.py 
223
224
225
226
227
228
229
230
231
232
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "BlockMatrix2x2":
    """Reconstruct this operator from JAX pytree children and static data."""
    del aux_data
    A, B, C, D = children
    return cls(A=A, B=B, C=C, D=D)

Toeplitz matrix operators with FFT-based multiplication.

Provides efficient O(n log n) matrix-vector products for Toeplitz matrices using FFT-based circular convolution.

Toeplitz

Bases: LinearOperator

Symmetric Toeplitz matrix represented by its first column/row v.

Source code in linox/operators/toeplitz.py 
 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
class Toeplitz(LinearOperator):
    """Symmetric Toeplitz matrix represented by its first column/row v."""

    def __init__(self, v: ArrayLike) -> None:
        self.v = jnp.asarray(v)
        n = int(self.v.shape[0])
        super().__init__((n, n), self.v.dtype)

    @property
    def shape(self) -> tuple[int, int]:
        """Shape of the Toeplitz matrix."""
        n = int(self.v.shape[0])
        return (n, n)

    def _matmul(self, vector: jax.Array) -> jax.Array:
        """Matrix-vector product for symmetric Toeplitz matrix.

        Uses FFT-based circulant convolution for O(n log n) complexity.

        Args:
            vector: Vector to multiply with, shape (n,) or (n, k)

        Returns
        -------
            Result of multiplication, shape (n,) or (n, k)
        """
        if vector.ndim == 1:
            vector = vector[:, None]
            squeeze_output = True
        else:
            squeeze_output = False

        # Input shape: (..., n, k)
        # We need to perform FFT along the dimension with size n (axis -2)
        n = self.shape[0]
        if vector.shape[-2] != n:
            msg = f"Dimension mismatch: expected size {n} at axis -2, got {vector.shape[-2]}"
            raise ValueError(msg)

        # Embed first row into circulant row: [v, 0, v[-1:0:-1]]
        # But for symmetric Toeplitz: [v, v[n-1:0:-1]] ?
        # Code was: jnp.concatenate([self.v, self.v[-1:0:-1]])
        # This creates the first column of the circulant matrix.
        embedded_col = jnp.concatenate([self.v, self.v[-1:0:-1]])
        p = embedded_col.shape[0]

        # FFT of the circulant column (1D)
        fft_col = jnp.fft.fft(embedded_col)

        # Pad vector along axis -2 to length p
        # padding shape: (..., p-n, k)
        pad_width = p - n
        padding_shape = list(vector.shape)
        padding_shape[-2] = pad_width
        zeros = jnp.zeros(padding_shape, dtype=vector.dtype)
        vector_padded = jnp.concatenate([vector, zeros], axis=-2)

        # FFT along axis -2
        fft_vector = jnp.fft.fft(vector_padded, axis=-2)

        # Broadcast fft_col: needs shape (1, ..., 1, p, 1) to match (..., p, k)
        # Actually standard broadcasting rules: (p,) broadcasts to (..., p, k) if p is last?
        # No, p is second to last.
        # We need (..., p, 1).
        # Reshape fft_col to (1, ..., 1, p, 1)
        # Simpler: expand dims at -1 and as many as needed on left
        fft_col.reshape((-1, 1))  # (p, 1)
        # This will broadcast against (..., p, k) correctly as (p, 1) * (..., p, k) ?
        # No. (p, 1) * (batch, p, k) -> (batch, p, k). Correct.
        # Wait, if batch is present. (B, p, k). (p, 1).
        # (p, 1) broadcasts to (B, p, k) ? No. Last dims must align.
        # (p, 1) aligns with (k,) ? No.
        # (B, p, k) * (p, 1) -> broadcasting (p, 1) against (k) fails.
        # We need to broadcast over k.
        # fft_vector: (..., p, k). fft_col: (p,).
        # We want fft_col to multiply along p axis.
        # Reshape fft_col to (p, 1).
        # (..., p, k) * (p, 1) -> (..., p, k).
        # Example: (2, 5, 3) * (5, 1).
        # (2, 5, 3). (5, 1).
        # 3 vs 1 -> ok.
        # 5 vs 5 -> ok.
        # 2 vs ? -> 1. ok.
        # So (p, 1) works!

        fft_result = fft_vector * fft_col.reshape((-1, 1))

        # IFFT
        result = jnp.fft.ifft(fft_result, axis=-2).real

        # Slice to original size n
        # This slicing doesn't work simply with [..., :n, :] syntax in python slices unless we build it?
        # Actually result[..., :n, :] syntax works in JAX/Numpy.
        result = result[..., :n, :]

        if squeeze_output:
            result = result.squeeze(axis=-1)  # Was axis=1 which assumes (n, 1) -> (n,)
            # If input was (n,), vector became (n, 1). Result (n, 1). Squeeze -> (n,).
            # If input was (1, n), vector became (1, n). Reshape?
            # Wait, line starts: if vector.ndim == 1: vector = vector[:, None].
            # (n,) -> (n, 1).
            # Squeeze -1 works.

        return result

    def _todense(self) -> jax.Array:
        return jsp.linalg.toeplitz(self.v)

    def from_matrix(self, matrix: jax.Array) -> "Toeplitz":
        """Create Toeplitz operator from matrix."""
        self.v = matrix[0, :]
        return Toeplitz(self.v)

    def transpose(self) -> "Toeplitz":
        """Return transpose (self for symmetric Toeplitz)."""
        return Toeplitz(self.v)

shape: tuple[int, int] property

Shape of the Toeplitz matrix.

from_matrix(matrix: jax.Array) -> Toeplitz

Create Toeplitz operator from matrix.

Source code in linox/operators/toeplitz.py 
133
134
135
136
def from_matrix(self, matrix: jax.Array) -> "Toeplitz":
    """Create Toeplitz operator from matrix."""
    self.v = matrix[0, :]
    return Toeplitz(self.v)

transpose() -> Toeplitz

Return transpose (self for symmetric Toeplitz).

Source code in linox/operators/toeplitz.py 
138
139
140
def transpose(self) -> "Toeplitz":
    """Return transpose (self for symmetric Toeplitz)."""
    return Toeplitz(self.v)

Permutation operations for linear operators.

This module implements permutation operations for linear operators, including:

  • :class:Permutation: Represents a permutation matrix :math:P that permutes the rows of a vector according to a given permutation

Permutation

Bases: LinearOperator

A linear operator defined via a permutation matrix.

For a permutation vector :math:p, this represents the permutation matrix :math:P where :math:P_{ij} = 1 if :math:j = p_i and :math:0 otherwise. The action on a vector :math:x is given by :math:(Px)_i = x_{p_i}, i.e., it permutes the elements of :math:x according to the permutation :math:p.

Args: perm: The permutation vector defining the operator perm_inv: The inverse permutation vector (optional, computed if not provided)

Source code in linox/operators/permutation.py 
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
class Permutation(LinearOperator):
    r"""A linear operator defined via a permutation matrix.

    For a permutation vector :math:`p`, this represents the permutation matrix :math:`P`
    where :math:`P_{ij} = 1` if :math:`j = p_i` and :math:`0` otherwise. The action on a
    vector :math:`x` is given by :math:`(Px)_i = x_{p_i}`, i.e., it permutes the
    elements of :math:`x` according to the permutation :math:`p`.

    Args:
        perm: The permutation vector defining the operator
        perm_inv: The inverse permutation vector (optional, computed if not provided)
    """

    def __init__(self, perm: ArrayLike, perm_inv: ArrayLike | None = None) -> None:
        self._perm = jnp.asarray(perm, dtype=jnp.int32)
        self._perm_inv = jnp.asarray(perm_inv, dtype=jnp.int32) if perm_inv is not None else jnp.argsort(self._perm, axis=-1)
        perm_size = self._perm.shape[-1]
        super().__init__(
            shape=(*self._perm.shape[:-1], perm_size, perm_size),
            dtype=default_floating_dtype(),  # Otherwise operation not allowed
        )

    # def _matmul(self, x: jnp.ndarray) -> jnp.ndarray:
    #     return _perm_op(x, self._perm)
    def _matmul(self, x: jax.Array) -> jax.Array:
        return _permute_rows(x, self._perm)

    def _todense(self) -> jax.Array:
        n = self.shape[-1]
        I = jnp.eye(n, dtype=self.dtype)
        return _permute_rows(I, self._perm)

    def transpose(self) -> "Permutation":
        """Return transposed permutation."""
        return Permutation(self._perm_inv, self._perm)

    def inverse(self) -> "Permutation":
        """Return inverse permutation."""
        return self.transpose()

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration."""
        children = (self._perm, self._perm_inv)
        aux_data = {}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "Permutation":
        """Unflatten for JAX pytree registration."""
        del aux_data
        perm, perm_inv = children
        return cls(perm=perm, perm_inv=perm_inv)

inverse() -> Permutation

Return inverse permutation.

Source code in linox/operators/permutation.py 
75
76
77
def inverse(self) -> "Permutation":
    """Return inverse permutation."""
    return self.transpose()

transpose() -> Permutation

Return transposed permutation.

Source code in linox/operators/permutation.py 
71
72
73
def transpose(self) -> "Permutation":
    """Return transposed permutation."""
    return Permutation(self._perm_inv, self._perm)

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

Source code in linox/operators/permutation.py 
79
80
81
82
83
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration."""
    children = (self._perm, self._perm_inv)
    aux_data = {}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> Permutation classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/permutation.py 
85
86
87
88
89
90
91
92
93
94
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "Permutation":
    """Unflatten for JAX pytree registration."""
    del aux_data
    perm, perm_inv = children
    return cls(perm=perm, perm_inv=perm_inv)

Triangular and Factor-based Linear Operators.

This module implements operators that are defined by their factors, such as: - :class:Triangular: A triangular matrix (L or U) - :class:CholeskyFactor: A specific lower triangular factor L from a Cholesky decomposition - :class:PSDFromFactor: A positive semi-definite operator defined as A = L @ L.T

CholeskyFactor

Bases: Triangular

A Cholesky factor L (lower triangular).

This operator represents the factor L such that A = L @ L.T. It specializes solves to use solve_triangular.

Source code in linox/operators/factor.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class CholeskyFactor(Triangular):
    """A Cholesky factor L (lower triangular).

    This operator represents the factor L such that A = L @ L.T.
    It specializes solves to use `solve_triangular`.
    """

    def __init__(self, L: jax.Array) -> None:
        # Cholesky factors are lower triangular by convention in Linox
        super().__init__(L, lower=True)

    def transpose(self) -> "Triangular":
        """Return the transpose of this operator."""
        # Transpose of CholeskyFactor is just a generic Upper Triangular matrix,
        # it loses the "CholeskyFactor" semantic meaning (which implies L is lower).
        return Triangular(self._A.T, lower=False)

transpose() -> Triangular

Return the transpose of this operator.

Source code in linox/operators/factor.py 
76
77
78
79
80
def transpose(self) -> "Triangular":
    """Return the transpose of this operator."""
    # Transpose of CholeskyFactor is just a generic Upper Triangular matrix,
    # it loses the "CholeskyFactor" semantic meaning (which implies L is lower).
    return Triangular(self._A.T, lower=False)

PSDFromFactor

Bases: LinearOperator

Positive Semi-Definite operator defined by a factor L: A = L @ L.T.

This is useful when we already have a matrix square root or Cholesky factor and want to represent the full covariance/operator without squaring it explicitly.

Source code in linox/operators/factor.py 
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
class PSDFromFactor(LinearOperator):
    """Positive Semi-Definite operator defined by a factor L: A = L @ L.T.

    This is useful when we already have a matrix square root or Cholesky factor
    and want to represent the full covariance/operator without squaring it explicitly.
    """

    def __init__(self, L: LinearOperator | jax.Array) -> None:
        if isinstance(L, (tuple, list)):
            # Handle tree unflatten case where internal might be passed incorrectly if not careful
            # But normal usage L is LinOp or Array
            pass

        self.L = utils.as_linop(L)
        shape = (self.L.shape[0], self.L.shape[0])
        super().__init__(shape, self.L.dtype)

    def _matmul(self, x: jax.Array) -> jax.Array:
        # A x = L (L.T x)
        return self.L @ (self.L.T @ x)

    def _todense(self) -> jax.Array:
        L_dense = self.L._todense()
        return L_dense @ L_dense.T

    def transpose(self) -> "PSDFromFactor":
        """Return the transpose of this operator."""
        return self

    def tree_flatten(self):
        """Flatten this operator into JAX pytree children and static data."""
        return (self.L,), {}

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        """Reconstruct this operator from JAX pytree children and static data."""
        return cls(children[0])

transpose() -> PSDFromFactor

Return the transpose of this operator.

Source code in linox/operators/factor.py 
129
130
131
def transpose(self) -> "PSDFromFactor":
    """Return the transpose of this operator."""
    return self

tree_flatten()

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/factor.py 
133
134
135
def tree_flatten(self):
    """Flatten this operator into JAX pytree children and static data."""
    return (self.L,), {}

tree_unflatten(aux_data, children) classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/factor.py 
137
138
139
140
@classmethod
def tree_unflatten(cls, aux_data, children):
    """Reconstruct this operator from JAX pytree children and static data."""
    return cls(children[0])

Triangular

Bases: LinearOperator

Triangular matrix operator.

Args: A: The triangular matrix data. lower: If True, A is lower triangular. If False, upper triangular.

Source code in linox/operators/factor.py 
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
class Triangular(LinearOperator):
    """Triangular matrix operator.

    Args:
        A: The triangular matrix data.
        lower: If True, A is lower triangular. If False, upper triangular.
    """

    def __init__(self, A: jax.Array, lower: bool = True) -> None:
        self._A = jnp.asarray(A)
        self._lower = lower
        super().__init__(self._A.shape, self._A.dtype)

    @property
    def lower(self) -> bool:
        """Whether the stored factor is lower triangular."""
        return self._lower

    def _matmul(self, x: jax.Array) -> jax.Array:
        # Triangular matmul is just standard matmul, but we know structure.
        # For now, just use dense matmul.
        return self._A @ x

    def _todense(self) -> jax.Array:
        if self._lower:
            return jnp.tril(self._A)
        return jnp.triu(self._A)

    def transpose(self) -> "Triangular":
        """Return the transpose of this operator."""
        return Triangular(self._A.T, lower=not self._lower)

    def tree_flatten(self):
        """Flatten this operator into JAX pytree children and static data."""
        return (self._A,), {"lower": self._lower}

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        """Reconstruct this operator from JAX pytree children and static data."""
        return cls(children[0], **aux_data)

lower: bool property

Whether the stored factor is lower triangular.

transpose() -> Triangular

Return the transpose of this operator.

Source code in linox/operators/factor.py 
51
52
53
def transpose(self) -> "Triangular":
    """Return the transpose of this operator."""
    return Triangular(self._A.T, lower=not self._lower)

tree_flatten()

Flatten this operator into JAX pytree children and static data.

Source code in linox/operators/factor.py 
55
56
57
def tree_flatten(self):
    """Flatten this operator into JAX pytree children and static data."""
    return (self._A,), {"lower": self._lower}

tree_unflatten(aux_data, children) classmethod

Reconstruct this operator from JAX pytree children and static data.

Source code in linox/operators/factor.py 
59
60
61
62
@classmethod
def tree_unflatten(cls, aux_data, children):
    """Reconstruct this operator from JAX pytree children and static data."""
    return cls(children[0], **aux_data)

_(A: PSDFromFactor) -> tuple[jax.Array, jax.Array]

log|A| = log|L L^T| = 2 log|L|.

Source code in linox/operators/factor.py 
203
204
205
206
207
208
209
210
@slogdet.dispatch
def _(A: PSDFromFactor) -> tuple[jax.Array, jax.Array]:
    """log|A| = log|L L^T| = 2 log|L|."""
    # We need slogdet of L.
    _sgn, logabs = slogdet(A.L)
    # A is PSD, so sign is 1.0 (unless L is singular/complex, but |A| >= 0)
    # log|A| = 2 * log|L|
    return jnp.array(1.0, dtype=A.dtype), 2 * logabs

Kernels

Kernel-based linear operators with automatic structure detection.

Provides lazy kernel operators that avoid materializing full kernel matrices, with automatic selection of Toeplitz structure for stationary kernels on uniform grids.

KernelOperator

Bases: LinearOperator

Base class for kernel-based linear operators.

Source code in linox/operators/kernel.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
class KernelOperator(LinearOperator):
    """Base class for kernel-based linear operators."""

    def __init__(
        self,
        kernel: Callable[[jax.Array, jax.Array], jax.Array],
        x0: jax.Array,
        x1: jax.Array | None = None,
    ) -> None:
        self.kernel = kernel
        self.x0 = x0
        self.x1 = x0 if x1 is None else x1
        super().__init__(shape=(self.x0.shape[0], self.x1.shape[0]), dtype=x0.dtype)

ArrayKernel

Bases: KernelOperator

Lazy kernel operator that NEVER materializes the full matrix.

All matrix-vector products are computed on-the-fly using lax.map for JIT compatibility. This operator is designed for large-scale problems where the full kernel matrix would exceed memory.

WARNING: Calling _todense() on large operators will cause OOM.

Args: kernel: Kernel function k(x, y) -> scalar x0: First set of points (n0, d) x1: Second set of points (n1, d), defaults to x0 if None chunk_size: Chunk size for chunked computation

Source code in linox/operators/kernel.py 
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
class ArrayKernel(KernelOperator):
    """Lazy kernel operator that NEVER materializes the full matrix.

    All matrix-vector products are computed on-the-fly using lax.map
    for JIT compatibility. This operator is designed for large-scale
    problems where the full kernel matrix would exceed memory.

    WARNING: Calling _todense() on large operators will cause OOM.

    Args:
        kernel: Kernel function k(x, y) -> scalar
        x0: First set of points (n0, d)
        x1: Second set of points (n1, d), defaults to x0 if None
        chunk_size: Chunk size for chunked computation
    """

    def __init__(
        self,
        kernel: Callable[[jax.Array, jax.Array], jax.Array],
        x0: jax.Array,
        x1: jax.Array | None = None,
        chunk_size: int = 256,
    ) -> None:
        super().__init__(kernel, x0, x1)
        self.chunk_size = chunk_size

    def _matmul(self, vec: jax.Array) -> jax.Array:
        """Matrix-free matmul using lax.map (JIT-compatible).

        Computes K @ v row-by-row without ever building K.
        Uses lax.map for JIT compatibility.
        """
        x0 = self.x0
        x1 = self.x1

        kernel_row_fn = jax.vmap(self.kernel, in_axes=(None, 0))

        def compute_row_dot(xi):
            row = kernel_row_fn(xi, x1)
            return jnp.dot(row, vec)

        return lax.map(compute_row_dot, x0)

    def transpose(self) -> "ArrayKernel":
        """Return transposed kernel operator."""
        return ArrayKernel(
            kernel=lambda x, y: self.kernel(y, x),
            x0=self.x1,
            x1=self.x0,
            chunk_size=self.chunk_size,
        )

    def _todense(self) -> jax.Array:
        """Materialize the full kernel matrix.

        WARNING: This will cause OOM for large operators.
        Only use for debugging with small n or when explicitly needed.
        """
        n0, n1 = self.shape
        if n0 > DENSE_THRESHOLD or n1 > DENSE_THRESHOLD:
            config.warn(f"Densifying large kernel ({n0}x{n1}). This may cause OOM. Consider using matrix-free operations instead.")

        x0 = self.x0
        x1 = self.x1

        kernel_fn = jax.vmap(
            jax.vmap(self.kernel, in_axes=(None, 0)),
            in_axes=(0, None),
        )

        return kernel_fn(x0, x1)

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration.

        The kernel function is static (placed in aux_data) so that the
        operator can be passed through ``jax.jit``.
        """
        children = (self.x0, self.x1)
        aux_data = {"kernel": self.kernel, "chunk_size": self.chunk_size}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "ArrayKernel":
        """Unflatten for JAX pytree registration."""
        x0, x1 = children
        return cls(kernel=aux_data["kernel"], x0=x0, x1=x1, chunk_size=aux_data["chunk_size"])

transpose() -> ArrayKernel

Return transposed kernel operator.

Source code in linox/operators/kernel.py 
162
163
164
165
166
167
168
169
def transpose(self) -> "ArrayKernel":
    """Return transposed kernel operator."""
    return ArrayKernel(
        kernel=lambda x, y: self.kernel(y, x),
        x0=self.x1,
        x1=self.x0,
        chunk_size=self.chunk_size,
    )

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

The kernel function is static (placed in aux_data) so that the operator can be passed through jax.jit.

Source code in linox/operators/kernel.py 
191
192
193
194
195
196
197
198
199
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration.

    The kernel function is static (placed in aux_data) so that the
    operator can be passed through ``jax.jit``.
    """
    children = (self.x0, self.x1)
    aux_data = {"kernel": self.kernel, "chunk_size": self.chunk_size}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> ArrayKernel classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/kernel.py 
201
202
203
204
205
206
207
208
209
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "ArrayKernel":
    """Unflatten for JAX pytree registration."""
    x0, x1 = children
    return cls(kernel=aux_data["kernel"], x0=x0, x1=x1, chunk_size=aux_data["chunk_size"])

ToeplitzKernel

Bases: KernelOperator

Kernel operator using Toeplitz structure for O(n log n) matmul via FFT.

This operator is completely matrix-free - it only stores the Toeplitz vector and uses FFT-based multiplication. Suitable for: - Self-covariance (x1 == x0) - Uniform 1D grid - Stationary kernel k(x,y) = f(x-y)

Note: Since we use symmetric Toeplitz matrices, transpose() returns self.

Args: kernel: Stationary kernel function x0: Points array (n,) or (n, 1) - must be uniform 1D grid chunk_size: Chunk size (for compatibility, not used in Toeplitz matmul)

Source code in linox/operators/kernel.py 
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
class ToeplitzKernel(KernelOperator):
    """Kernel operator using Toeplitz structure for O(n log n) matmul via FFT.

    This operator is completely matrix-free - it only stores the Toeplitz vector
    and uses FFT-based multiplication. Suitable for:
    - Self-covariance (x1 == x0)
    - Uniform 1D grid
    - Stationary kernel k(x,y) = f(x-y)

    Note: Since we use symmetric Toeplitz matrices, transpose() returns self.

    Args:
        kernel: Stationary kernel function
        x0: Points array (n,) or (n, 1) - must be uniform 1D grid
        chunk_size: Chunk size (for compatibility, not used in Toeplitz matmul)
    """

    def __init__(
        self,
        kernel: Callable[[jax.Array, jax.Array], jax.Array],
        x0: jax.Array,
        x1: jax.Array | None = None,
        chunk_size: int = 256,
    ) -> None:
        if x1 is not None and not _is_self_covariance_cheap(x0, x1):
            msg = "ToeplitzKernel requires self-covariance (x1 must be None or x1 is x0)"
            raise ValueError(msg)

        super().__init__(kernel, x0, None)
        self.chunk_size = chunk_size
        self._toeplitz_vector = self._compute_toeplitz_vector()
        self._toeplitz_op = Toeplitz(self._toeplitz_vector)

    def _compute_toeplitz_vector(self) -> jax.Array:
        """Compute the first row of the Toeplitz matrix."""
        x0 = self.x0.ravel() if self.x0.ndim == 2 else self.x0
        x0_reshaped = x0.reshape(-1, 1) if x0.ndim == 1 else x0
        return jax.vmap(lambda xi: self.kernel(x0_reshaped[0], xi))(x0_reshaped)

    def _matmul(self, vec: jax.Array) -> jax.Array:
        """FFT-based O(n log n) matmul - completely matrix-free."""
        return self._toeplitz_op @ vec

    def transpose(self) -> "ToeplitzKernel":
        """Symmetric Toeplitz: transpose is self."""
        return self

    def _todense(self) -> jax.Array:
        """Materialize via Toeplitz operator.

        WARNING: This defeats the purpose of using ToeplitzKernel.
        """
        n = self.shape[0]
        if n > DENSE_THRESHOLD:
            config.warn(f"Densifying large ToeplitzKernel ({n}x{n}). This may cause OOM.")
        return self._toeplitz_op._todense()

    def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
        """Flatten for JAX pytree registration.

        The kernel function is static (placed in aux_data) so that the
        operator can be passed through ``jax.jit``.
        """
        children = (self.x0,)
        aux_data = {"kernel": self.kernel, "chunk_size": self.chunk_size}
        return children, aux_data

    @classmethod
    def tree_unflatten(
        cls,
        aux_data: dict[str, any],
        children: tuple[any, ...],
    ) -> "ToeplitzKernel":
        """Unflatten for JAX pytree registration."""
        (x0,) = children
        return cls(kernel=aux_data["kernel"], x0=x0, chunk_size=aux_data["chunk_size"])

transpose() -> ToeplitzKernel

Symmetric Toeplitz: transpose is self.

Source code in linox/operators/kernel.py 
264
265
266
def transpose(self) -> "ToeplitzKernel":
    """Symmetric Toeplitz: transpose is self."""
    return self

tree_flatten() -> tuple[tuple[any, ...], dict[str, any]]

Flatten for JAX pytree registration.

The kernel function is static (placed in aux_data) so that the operator can be passed through jax.jit.

Source code in linox/operators/kernel.py 
278
279
280
281
282
283
284
285
286
def tree_flatten(self) -> tuple[tuple[any, ...], dict[str, any]]:
    """Flatten for JAX pytree registration.

    The kernel function is static (placed in aux_data) so that the
    operator can be passed through ``jax.jit``.
    """
    children = (self.x0,)
    aux_data = {"kernel": self.kernel, "chunk_size": self.chunk_size}
    return children, aux_data

tree_unflatten(aux_data: dict[str, any], children: tuple[any, ...]) -> ToeplitzKernel classmethod

Unflatten for JAX pytree registration.

Source code in linox/operators/kernel.py 
288
289
290
291
292
293
294
295
296
@classmethod
def tree_unflatten(
    cls,
    aux_data: dict[str, any],
    children: tuple[any, ...],
) -> "ToeplitzKernel":
    """Unflatten for JAX pytree registration."""
    (x0,) = children
    return cls(kernel=aux_data["kernel"], x0=x0, chunk_size=aux_data["chunk_size"])

kernel_operator(kernel: Callable[[jax.Array, jax.Array], jax.Array], x0: jax.Array, x1: jax.Array | None = None, is_stationary: bool = False, assume_uniform: bool = False, chunk_size: int = 256) -> KernelOperator

Create the optimal kernel operator based on structure detection.

Automatically selects ToeplitzKernel when: - Self-covariance (x1 is None or x1 is x0 - identity check only, no value comparison) - Points are uniform 1D grid (assume_uniform=True or host-side check for small n) - Kernel is stationary (is_stationary=True)

Otherwise creates a lazy ArrayKernel that never materializes the full matrix.

Parameters:

Name Type Description Default
kernel Callable[[Array, Array], Array]

Kernel function k(x, y) -> scalar

 required
x0 Array

First set of points

 required
x1 Array

Second set of points (None for self-covariance)

 None
is_stationary bool

True if kernel is stationary k(x,y) = f(x-y)

 False
assume_uniform bool

True to skip uniformity check (use when creating grid with arange)

 False
chunk_size int

Chunk size for lazy matmul computation

 256

Returns:

Type Description
KernelOperator

Either ToeplitzKernel or ArrayKernel depending on structure
Source code in linox/operators/kernel.py 
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
def kernel_operator(
    kernel: Callable[[jax.Array, jax.Array], jax.Array],
    x0: jax.Array,
    x1: jax.Array | None = None,
    is_stationary: bool = False,
    assume_uniform: bool = False,
    chunk_size: int = 256,
) -> "KernelOperator":
    """Create the optimal kernel operator based on structure detection.

    Automatically selects ToeplitzKernel when:
    - Self-covariance (x1 is None or x1 is x0 - identity check only, no value comparison)
    - Points are uniform 1D grid (assume_uniform=True or host-side check for small n)
    - Kernel is stationary (is_stationary=True)

    Otherwise creates a lazy ArrayKernel that never materializes the full matrix.

    Parameters
    ----------
    kernel : Callable[[jax.Array, jax.Array], jax.Array]
        Kernel function k(x, y) -> scalar
    x0 : jax.Array
        First set of points
    x1 : jax.Array, optional
        Second set of points (None for self-covariance)
    is_stationary : bool, default=False
        True if kernel is stationary k(x,y) = f(x-y)
    assume_uniform : bool, default=False
        True to skip uniformity check (use when creating grid with arange)
    chunk_size : int, default=256
        Chunk size for lazy matmul computation

    Returns
    -------
    KernelOperator
        Either ToeplitzKernel or ArrayKernel depending on structure
    """
    is_self_cov = _is_self_covariance_cheap(x0, x1)
    is_uniform = assume_uniform or _is_uniform_1d_host(x0)

    if is_self_cov and is_uniform and is_stationary:
        return ToeplitzKernel(kernel, x0, chunk_size=chunk_size)

    return ArrayKernel(kernel, x0, x1, chunk_size=chunk_size)

Property wrappers

Property wrapper operators for symmetry and PSD assumptions.

This module provides lightweight wrapper operators that tag operands with semantic properties like symmetry (Sym) and positive semidefiniteness (PSD).

These wrappers enable specialized algorithm dispatch without modifying the underlying operator. They are compositional and may propagate under arithmetic operations.

See Also

ADR-0003 : Architecture decision record for PSD and symmetry wrappers.

Examples:

>>> import linox as lo
>>> K = lo.Diagonal(jnp.array([1.0, 2.0, 3.0]))
>>> K_psd = lo.PSD(K)  # Promise that K is PSD
>>> K_psd.is_psd
True

PSD

Bases: LinearOperator

Wrapper declaring an operator as positive semidefinite.

This is a semantic promise that x^T A x >= 0 for all x. The wrapper does not verify this numerically unless debug validation is performed.

PSD operators are also implicitly symmetric.

Parameters:

Name Type Description Default
op LinearOperator

The operator to wrap. Must be square.

 required

Attributes:

Name Type Description
is_psd bool

Always True for PSD-wrapped operators.
is_symmetric bool

Always True (PSD implies symmetric).
wrapped LinearOperator

The underlying wrapped operator.

Raises:

Type Description
ValueError

If the operator is not square.
See Also

Sym : Wrapper for symmetric operators. assume_psd : Convenience function to create PSD wrappers.

Examples:

>>> K = lo.Diagonal(jnp.array([1.0, 2.0, 3.0]))
>>> K_psd = lo.PSD(K)
>>> K_psd.is_psd
True
>>> K_psd.is_symmetric
True
Source code in linox/operators/wrappers.py 
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
class PSD(LinearOperator):
    """Wrapper declaring an operator as positive semidefinite.

    This is a semantic promise that x^T A x >= 0 for all x. The wrapper
    does not verify this numerically unless debug validation is performed.

    PSD operators are also implicitly symmetric.

    Parameters
    ----------
    op : LinearOperator
        The operator to wrap. Must be square.

    Attributes
    ----------
    is_psd : bool
        Always True for PSD-wrapped operators.
    is_symmetric : bool
        Always True (PSD implies symmetric).
    wrapped : LinearOperator
        The underlying wrapped operator.

    Raises
    ------
    ValueError
        If the operator is not square.

    See Also
    --------
    Sym : Wrapper for symmetric operators.
    assume_psd : Convenience function to create PSD wrappers.

    Examples
    --------
    >>> K = lo.Diagonal(jnp.array([1.0, 2.0, 3.0]))
    >>> K_psd = lo.PSD(K)
    >>> K_psd.is_psd
    True
    >>> K_psd.is_symmetric
    True
    """

    wrapped: LinearOperator

    def __init__(self, op: LinearOperator) -> None:
        self.wrapped = op

        if op.shape[-2] != op.shape[-1]:
            msg = f"PSD wrapper requires square operator, got shape {op.shape}"
            raise ValueError(msg)

        super().__init__(shape=op.shape, dtype=op.dtype)

    @property
    def is_symmetric(self) -> bool:
        """Return True, since PSD operators are symmetric."""
        return True

    @property
    def is_psd(self) -> bool:
        """Return True, indicating this operator is declared PSD."""
        return True

    def _matmul(self, other: jax.Array) -> jax.Array:
        """Compute matrix-matrix product by delegating to wrapped operator."""
        return self.wrapped @ other

    def transpose(self) -> LinearOperator:
        """Return self, since PSD operators are symmetric."""
        return self

    def _todense(self) -> jax.Array:
        """Return dense representation of the wrapped operator."""
        return self.wrapped.todense()

    def children(self) -> tuple[LinearOperator, ...]:
        """Return child operators for tree traversal."""
        return (self.wrapped,)

    def tree_flatten(self) -> tuple[tuple, dict]:
        """Flatten for JAX PyTree compatibility."""
        return (self.wrapped,), {}

    @classmethod
    def tree_unflatten(cls, aux: dict, children: tuple) -> PSD:
        """Unflatten from JAX PyTree representation."""
        (wrapped,) = children
        return cls(wrapped)

is_psd: bool property

Return True, indicating this operator is declared PSD.

is_symmetric: bool property

Return True, since PSD operators are symmetric.

children() -> tuple[LinearOperator, ...]

Return child operators for tree traversal.

Source code in linox/operators/wrappers.py 
196
197
198
def children(self) -> tuple[LinearOperator, ...]:
    """Return child operators for tree traversal."""
    return (self.wrapped,)

transpose() -> LinearOperator

Return self, since PSD operators are symmetric.

Source code in linox/operators/wrappers.py 
188
189
190
def transpose(self) -> LinearOperator:
    """Return self, since PSD operators are symmetric."""
    return self

tree_flatten() -> tuple[tuple, dict]

Flatten for JAX PyTree compatibility.

Source code in linox/operators/wrappers.py 
200
201
202
def tree_flatten(self) -> tuple[tuple, dict]:
    """Flatten for JAX PyTree compatibility."""
    return (self.wrapped,), {}

tree_unflatten(aux: dict, children: tuple) -> PSD classmethod

Unflatten from JAX PyTree representation.

Source code in linox/operators/wrappers.py 
204
205
206
207
208
@classmethod
def tree_unflatten(cls, aux: dict, children: tuple) -> PSD:
    """Unflatten from JAX PyTree representation."""
    (wrapped,) = children
    return cls(wrapped)

SPD

Bases: PSD

Wrapper declaring an operator as symmetric positive definite.

This is a semantic promise that x^T A x > 0 for all x != 0 (strictly positive definite). SPD operators are a subset of PSD operators.

Parameters:

Name Type Description Default
op LinearOperator

The operator to wrap. Must be square.

 required

Attributes:

Name Type Description
is_spd bool

Always True for SPD-wrapped operators.
is_psd bool

Always True (SPD implies PSD).
is_symmetric bool

Always True (SPD implies symmetric).
See Also

PSD : Wrapper for positive semidefinite operators. Sym : Wrapper for symmetric operators.

Examples:

>>> K = lo.Diagonal(jnp.array([1.0, 2.0, 3.0]))  # All positive
>>> K_spd = lo.SPD(K)
>>> K_spd.is_spd
True
Source code in linox/operators/wrappers.py 
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
class SPD(PSD):
    """Wrapper declaring an operator as symmetric positive definite.

    This is a semantic promise that x^T A x > 0 for all x != 0 (strictly
    positive definite). SPD operators are a subset of PSD operators.

    Parameters
    ----------
    op : LinearOperator
        The operator to wrap. Must be square.

    Attributes
    ----------
    is_spd : bool
        Always True for SPD-wrapped operators.
    is_psd : bool
        Always True (SPD implies PSD).
    is_symmetric : bool
        Always True (SPD implies symmetric).

    See Also
    --------
    PSD : Wrapper for positive semidefinite operators.
    Sym : Wrapper for symmetric operators.

    Examples
    --------
    >>> K = lo.Diagonal(jnp.array([1.0, 2.0, 3.0]))  # All positive
    >>> K_spd = lo.SPD(K)
    >>> K_spd.is_spd
    True
    """

    @property
    def is_spd(self) -> bool:
        """Return True, indicating this operator is declared SPD."""
        return True

    @classmethod
    def tree_unflatten(cls, aux: dict, children: tuple) -> SPD:
        """Unflatten from JAX PyTree representation."""
        (wrapped,) = children
        return cls(wrapped)

is_spd: bool property

Return True, indicating this operator is declared SPD.

tree_unflatten(aux: dict, children: tuple) -> SPD classmethod

Unflatten from JAX PyTree representation.

Source code in linox/operators/wrappers.py 
249
250
251
252
253
@classmethod
def tree_unflatten(cls, aux: dict, children: tuple) -> SPD:
    """Unflatten from JAX PyTree representation."""
    (wrapped,) = children
    return cls(wrapped)

Sym

Bases: LinearOperator

Wrapper declaring an operator as symmetric (self-adjoint).

This is a semantic promise that A = A^T (or A = A^H for complex). The wrapper does not verify this numerically unless debug validation is performed.

Parameters:

Name Type Description Default
op LinearOperator

The operator to wrap. Must be square.

 required

Attributes:

Name Type Description
is_symmetric bool

Always True for Sym-wrapped operators.
wrapped LinearOperator

The underlying wrapped operator.

Raises:

Type Description
ValueError

If the operator is not square.
See Also

PSD : Wrapper for positive semidefinite operators. assume_symmetric : Convenience function to create Sym wrappers.

Examples:

>>> A = lo.Matrix(jnp.array([[1, 2], [2, 1]]))
>>> A_sym = lo.Sym(A)
>>> A_sym.is_symmetric
True
Source code in linox/operators/wrappers.py 
 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
class Sym(LinearOperator):
    """Wrapper declaring an operator as symmetric (self-adjoint).

    This is a semantic promise that A = A^T (or A = A^H for complex).
    The wrapper does not verify this numerically unless debug validation
    is performed.

    Parameters
    ----------
    op : LinearOperator
        The operator to wrap. Must be square.

    Attributes
    ----------
    is_symmetric : bool
        Always True for Sym-wrapped operators.
    wrapped : LinearOperator
        The underlying wrapped operator.

    Raises
    ------
    ValueError
        If the operator is not square.

    See Also
    --------
    PSD : Wrapper for positive semidefinite operators.
    assume_symmetric : Convenience function to create Sym wrappers.

    Examples
    --------
    >>> A = lo.Matrix(jnp.array([[1, 2], [2, 1]]))
    >>> A_sym = lo.Sym(A)
    >>> A_sym.is_symmetric
    True
    """

    wrapped: LinearOperator

    def __init__(self, op: LinearOperator) -> None:
        self.wrapped = op

        if op.shape[-2] != op.shape[-1]:
            msg = f"Sym wrapper requires square operator, got shape {op.shape}"
            raise ValueError(msg)

        super().__init__(shape=op.shape, dtype=op.dtype)

    @property
    def is_symmetric(self) -> bool:
        """Return True, indicating this operator is declared symmetric."""
        return True

    @property
    def is_psd(self) -> bool:
        """Return the PSD status of the wrapped operator."""
        return getattr(self.wrapped, "is_psd", False)

    def _matmul(self, other: jax.Array) -> jax.Array:
        """Compute matrix-matrix product by delegating to wrapped operator."""
        return self.wrapped @ other

    def transpose(self) -> LinearOperator:
        """Return self, since symmetric operators equal their transpose."""
        return self

    def _todense(self) -> jax.Array:
        """Return dense representation of the wrapped operator."""
        return self.wrapped.todense()

    def children(self) -> tuple[LinearOperator, ...]:
        """Return child operators for tree traversal."""
        return (self.wrapped,)

    def tree_flatten(self) -> tuple[tuple, dict]:
        """Flatten for JAX PyTree compatibility."""
        return (self.wrapped,), {}

    @classmethod
    def tree_unflatten(cls, aux: dict, children: tuple) -> Sym:
        """Unflatten from JAX PyTree representation."""
        (wrapped,) = children
        return cls(wrapped)

is_psd: bool property

Return the PSD status of the wrapped operator.

is_symmetric: bool property

Return True, indicating this operator is declared symmetric.

children() -> tuple[LinearOperator, ...]

Return child operators for tree traversal.

Source code in linox/operators/wrappers.py 
106
107
108
def children(self) -> tuple[LinearOperator, ...]:
    """Return child operators for tree traversal."""
    return (self.wrapped,)

transpose() -> LinearOperator

Return self, since symmetric operators equal their transpose.

Source code in linox/operators/wrappers.py 
 98
 99
100
def transpose(self) -> LinearOperator:
    """Return self, since symmetric operators equal their transpose."""
    return self

tree_flatten() -> tuple[tuple, dict]

Flatten for JAX PyTree compatibility.

Source code in linox/operators/wrappers.py 
110
111
112
def tree_flatten(self) -> tuple[tuple, dict]:
    """Flatten for JAX PyTree compatibility."""
    return (self.wrapped,), {}

tree_unflatten(aux: dict, children: tuple) -> Sym classmethod

Unflatten from JAX PyTree representation.

Source code in linox/operators/wrappers.py 
114
115
116
117
118
@classmethod
def tree_unflatten(cls, aux: dict, children: tuple) -> Sym:
    """Unflatten from JAX PyTree representation."""
    (wrapped,) = children
    return cls(wrapped)

_(a: PSD) -> LinearOperator

Delegate to the wrapped operator. Also covers SPD, which extends PSD.

Source code in linox/operators/wrappers.py 
370
371
372
373
@lsqrt.dispatch
def _(a: PSD) -> LinearOperator:
    """Delegate to the wrapped operator. Also covers `SPD`, which extends `PSD`."""
    return lsqrt(a.wrapped)

assume_psd(op: LinearOperator) -> PSD

Wrap an operator to declare it as positive semidefinite.

Parameters:

Name Type Description Default
op LinearOperator

The operator to wrap. Must be square.

 required

Returns:

Type Description
PSD

The wrapped operator with is_psd=True and is_symmetric=True.
See Also

PSD : The wrapper class. assume_symmetric : For symmetric operators without PSD guarantee.

Examples:

>>> K = lo.Diagonal(jnp.array([1.0, 2.0, 3.0]))
>>> K_psd = lo.assume_psd(K)
>>> K_psd.is_psd
True
Source code in linox/operators/wrappers.py 
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
def assume_psd(op: LinearOperator) -> PSD:
    """Wrap an operator to declare it as positive semidefinite.

    Parameters
    ----------
    op : LinearOperator
        The operator to wrap. Must be square.

    Returns
    -------
    PSD
        The wrapped operator with is_psd=True and is_symmetric=True.

    See Also
    --------
    PSD : The wrapper class.
    assume_symmetric : For symmetric operators without PSD guarantee.

    Examples
    --------
    >>> K = lo.Diagonal(jnp.array([1.0, 2.0, 3.0]))
    >>> K_psd = lo.assume_psd(K)
    >>> K_psd.is_psd
    True
    """
    if isinstance(op, PSD):
        return op
    return PSD(op)

assume_spd(op: LinearOperator) -> SPD

Wrap an operator to declare it as symmetric positive definite.

Parameters:

Name Type Description Default
op LinearOperator

The operator to wrap. Must be square.

 required

Returns:

Type Description
SPD

The wrapped operator with is_spd=True, is_psd=True, is_symmetric=True.
See Also

SPD : The wrapper class. assume_psd : For positive semidefinite operators.

Examples:

>>> K = lo.Diagonal(jnp.array([1.0, 2.0, 3.0]))
>>> K_spd = lo.assume_spd(K)
>>> K_spd.is_spd
True
Source code in linox/operators/wrappers.py 
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
def assume_spd(op: LinearOperator) -> SPD:
    """Wrap an operator to declare it as symmetric positive definite.

    Parameters
    ----------
    op : LinearOperator
        The operator to wrap. Must be square.

    Returns
    -------
    SPD
        The wrapped operator with is_spd=True, is_psd=True, is_symmetric=True.

    See Also
    --------
    SPD : The wrapper class.
    assume_psd : For positive semidefinite operators.

    Examples
    --------
    >>> K = lo.Diagonal(jnp.array([1.0, 2.0, 3.0]))
    >>> K_spd = lo.assume_spd(K)
    >>> K_spd.is_spd
    True
    """
    if isinstance(op, SPD):
        return op
    return SPD(op)

assume_symmetric(op: LinearOperator) -> Sym

Wrap an operator to declare it as symmetric.

Parameters:

Name Type Description Default
op LinearOperator

The operator to wrap. Must be square.

 required

Returns:

Type Description
Sym

The wrapped operator with is_symmetric=True.
See Also

Sym : The wrapper class. assume_psd : For positive semidefinite operators.

Examples:

>>> A = lo.Matrix(jnp.array([[1, 2], [2, 1]]))
>>> A_sym = lo.assume_symmetric(A)
>>> A_sym.is_symmetric
True
Source code in linox/operators/wrappers.py 
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
def assume_symmetric(op: LinearOperator) -> Sym:
    """Wrap an operator to declare it as symmetric.

    Parameters
    ----------
    op : LinearOperator
        The operator to wrap. Must be square.

    Returns
    -------
    Sym
        The wrapped operator with is_symmetric=True.

    See Also
    --------
    Sym : The wrapper class.
    assume_psd : For positive semidefinite operators.

    Examples
    --------
    >>> A = lo.Matrix(jnp.array([[1, 2], [2, 1]]))
    >>> A_sym = lo.assume_symmetric(A)
    >>> A_sym.is_symmetric
    True
    """
    if isinstance(op, Sym):
        return op
    return Sym(op)