Linear operators

The question this method answers¶

How can code carry a matrix-like transformation without materializing every sum, product, transpose, or block as one dense array? A linear operator records how a map acts on a vector and how its transpose acts on a cotangent.

Before computation: what should be true?¶

Name the domain and codomain dimensions. An operator with shape $(m,n)$ accepts a vector of shape $(n,)$ in matvec and a vector of shape $(m,)$ in rmatvec. Composition requires matching inner dimensions, addition requires identical shapes, and a block diagonal operator requires at least one block.

Define the mathematical objects¶

A map $A$ is linear when scalar coefficients distribute over inputs:

The transpose action $A^\mathsf{T}y$ maps a cotangent in the output space back to the input space. A PyTree is a nested JAX structure with array leaves and static structure. Jaxstro operators are Equinox modules: matrix, diagonal, and scalar values are differentiable PyTree leaves, while the chosen composition tree and its shapes are static program structure.

Derive the method¶

The defining linearity relation is

For compatible maps $B:\mathbb{R}^n\rightarrow\mathbb{R}^k$ and $A:\mathbb{R}^k\rightarrow\mathbb{R}^m$, composition applies the right map first:

Reverse multiplication follows by reversing the factors. The defining adjoint identity is

These relations give independent audits that do not depend on the operator’s internal representation.

What the algorithm actually does¶

DenseOperator stores an $(m,n)$ array. DiagonalOperator stores an $(n,)$ diagonal. scale, add, compose, and transpose wrap existing operators; they do not materialize a matrix during matvec or rmatvec. block_diag(*blocks) slices an input according to block widths, applies each block, and concatenates the outputs. to_dense() is deliberately public so small and moderate fixtures can be compared against explicit matrix algebra.

The block implementation uses Python loops over the static block tuple. Python structure is static during tracing; the array work inside each operator remains JAX-traceable. Changing the number or types of blocks creates a different program structure and can trigger recompilation.

Thus shape checks happen eagerly for composition metadata; vector shape errors remain ordinary JAX array-algebra errors at application time.

What JAX differentiates¶

JAX differentiates matvec, rmatvec, and to_dense with respect to floating array leaves and floating input vectors. For a scalar loss built from $Ax$, AD computes derivatives of the executed composition, including sensitivities to a dense matrix, a diagonal, or a scale leaf. It does not differentiate the Protocol, tuple length, shapes, slicing offsets, or choice of operator class.

An operator transpose is an algebraic transpose, not a custom derivative rule. No solve occurs, so these classes make no implicit-solution derivative claim.

Using it in Jaxstro¶

import jax
import jax.numpy as jnp

from jaxstro.numerics.operators import DenseOperator, DiagonalOperator, compose

left_matrix = jnp.array([[1.0, 2.0], [0.0, 1.0]])
right_matrix = jnp.array([[2.0, 0.0], [1.0, 3.0]])
left = DenseOperator(left_matrix)
right = DiagonalOperator(jnp.array([2.0, 3.0]))
operator = compose(left, right)

x = jnp.array([0.5, -1.0])
y = jnp.array([1.5, 0.25])
forward = operator.matvec(x)
reverse = operator.rmatvec(y)
dense = operator.to_dense()
gradient = jax.grad(lambda vector: jnp.sum(operator.matvec(vector) ** 2))(x)

assert operator.shape == (2, 2)
assert jnp.allclose(forward, dense @ x)
assert jnp.allclose(reverse, dense.T @ y)
assert jnp.allclose(jnp.vdot(y, forward), jnp.vdot(reverse, x))
assert gradient.shape == x.shape

How to audit the result¶

1. Compare matvec(x) with to_dense() @ x on several nontrivial vectors.

2. Compare rmatvec(y) with to_dense().T @ y.

3. Check the adjoint inner-product identity to float tolerance.

4. Compare sums, products, transposes, and block assembly with explicit matrices.

5. For claimed leaf gradients, compare AD with central finite differences while the composition tree and shapes remain fixed.

Where the claim stops¶

The module does not provide sparse formats, iterative solves, preconditioners, shape-polymorphic block structure, or custom implicit differentiation. Dense parity demonstrates that a represented operator matches its explicit matrix on the tested fixtures; it does not establish conditioning or solver convergence.

Connected ideas¶