Tutorial notebook. Source .ipynb on GitHub.

Linear Operators and the Kalman Filter¶

This notebook shows how LinearOperator avoids materializing large matrices in the Kalman filter, enabling state dimensions of 10,000+ without hitting memory limits.

Motivation¶

The standard Kalman filter stores transition and observation matrices as dense arrays. For a state dimension d, the transition matrix is d x d — that's 80 MB at d=3000 in float64, 800 MB at d=10,000.

A matrix-free Function operator wraps a function v -> A @ v instead. It uses O(d) memory (just the closure) and supports @, +, .T and composition — the same algebra as dense matrices, without materialization.

import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)

from probjax.utils.linear_operator import Diagonal, Function

1. Creating a LinearOperator¶

A Function operator wraps any linear function v -> A @ v plus its input shape; the output shape and dtype are inferred without running the function. The @partial(Function, in_shape=...) decorator is the convenience syntax. Structured operators keep their structure: Diagonal(d) stores only the diagonal and Dense(m) a materialized matrix.

from functools import partial

# Tridiagonal matrix: 2*I + tril(ones) — O(d) memory, never materialized
d = 100

@partial(Function, in_shape=d)
def tridiag(x):
    return 2.0 * x - 0.5 * jnp.roll(x, 1) - 0.5 * jnp.roll(x, -1)

# Apply to a vector — no matrix stored
v = jnp.ones(d)
print("tridiag @ v shape:", (tridiag @ v).shape)
print("tridiag @ v first 5:", (tridiag @ v)[:5])

tridiag @ v shape: (100,)
tridiag @ v first 5: [1. 1. 1. 1. 1.]

2. Operators compose like matrices¶

You can do A @ B, A + B, A.T, A @ v — all without materializing.

# Diagonal operator
diag_vals = jnp.logspace(-1, 1, d)  # condition number 100
Diag = Diagonal(diag_vals)

# Compose: Diag @ tridiag
composed = Diag @ tridiag
print("Composed @ v first 5:", (composed @ v)[:5])

# Add: Diag + tridiag
added = Diag + tridiag
print("Added @ v first 5:", (added @ v)[:5])

# Transpose
print("Transpose @ v first 5:", (tridiag.T @ v)[:5])

Composed @ v first 5: [0.1        0.10476158 0.10974988 0.1149757  0.12045035]
Added @ v first 5: [1.1        1.10476158 1.10974988 1.1149757  1.12045035]
Transpose @ v first 5: [1. 1. 1. 1. 1.]

3. Materialize when needed¶

Call .as_array() to get the dense matrix. This is O(d²) in memory and calls the function d times.

mat = tridiag.as_array()
print("Dense shape:", mat.shape)
print("Memory: {:.1f} MB".format(mat.nbytes / 1024**2))
print("Is tridiagonal:", jnp.allclose(mat, mat.T))

Dense shape: (100, 100)
Memory: 0.1 MB
Is tridiagonal: True

4. Kalman filter with LinearOperators¶

The kalman_filter accepts LinearOperator for transition and observation matrices. The filter uses matrix-free operations internally (CG for solve, Lanczos for logdet) when dimensions are large, and materializes + dense Cholesky factorization when dimensions are small.

from probjax.inference.filtering.kalman_filter import kalman_filter

d = 1000
mu0 = jnp.ones(d)
C0 = jnp.eye(d) * 0.1

# Transition: tridiagonal (never materialized)
@partial(Function, in_shape=d)
def transition_fn(x):
    lower = -0.01 * x[:-1]
    diag = 0.99 * x
    return diag + jnp.pad(lower, (1, 0))

# Observation: only observe first 2 components
@partial(Function, in_shape=d)
def observation_fn(x):
    return jnp.array([x[0], x[1]])

# Process noise: diagonal (O(1) memory)
Q = Diagonal(jnp.asarray(0.001), d)

# No observation noise
def transition_model(t_old, t):
    return transition_fn, Q

def observation_model(t):
    return observation_fn, None

kernel = kalman_filter(transition_model, observation_model)
state = kernel.init(mu0, C0)

import time

# JIT and warmup
step_fn = jax.jit(lambda s: kernel.step(s, observed=jnp.array([1.0, 1.0])))
_ = step_fn(state)

# Benchmark
t0 = time.perf_counter()
for _ in range(10):
    new_state, _ = step_fn(state)
    new_state.mean.block_until_ready()
t_ms = (time.perf_counter() - t0) / 10 * 1000
print(f"State dim {d}: {t_ms:.1f} ms/step")
print(f"Memory: {d*d*8/1024**2:.0f} MB if materialized (not stored)")

State dim 1000: 25.3 ms/step
Memory: 8 MB if materialized (not stored)

5. Dense vs LinearOperator comparison¶

For small d, dense is faster (direct LAPACK factorization). For large d, LinearOperators win on memory and eventually on speed too.

# Dense version of the same model
transition_dense = transition_fn.as_array()
Q_dense = Q.as_array()
obs_dense = observation_fn.as_array()

def transition_model_dense(t_old, t):
    return transition_dense, Q_dense

def observation_model_dense(t):
    return obs_dense, None

kernel_dense = kalman_filter(transition_model_dense, observation_model_dense)
state_dense = kernel_dense.init(mu0, C0)

# Benchmark dense
step_dense = jax.jit(lambda s: kernel_dense.step(s, observed=jnp.array([1.0, 1.0])))
_ = step_dense(state_dense)  # warmup

t0 = time.perf_counter()
for _ in range(10):
    new_state_dense, _ = step_dense(state_dense)
    new_state_dense.mean.block_until_ready()
t_dense = (time.perf_counter() - t0) / 10 * 1000

print(f"d={d}  Dense: {t_dense:.1f} ms  |  LinearOp: {t_ms:.1f} ms")
print(f"Dense memory: {d*d*8/1024**2:.0f} MB  |  LinearOp: ~0 MB (not stored)")

d=1000  Dense: 29.3 ms  |  LinearOp: 25.3 ms
Dense memory: 8 MB  |  LinearOp: ~0 MB (not stored)

6. Scaling with dimension¶

Let's sweep over dimensions to see where LinearOperators start winning.

print(f"{'dim':>6} | {'dense (ms)':>10} | {'LO (ms)':>10} | {'mem (MB)':>8}")
print("-" * 45)

for dim in [100, 500, 1000, 2000, 3000]:
    mu = jnp.ones(dim)
    cov = jnp.eye(dim) * 0.1

    @partial(Function, in_shape=dim)
    def tr(x):
        return 0.99 * x - 0.01 * jnp.roll(x, 1)

    Q_lo = Diagonal(jnp.asarray(0.001), dim)

    @partial(Function, in_shape=dim)
    def obs(x):
        return jnp.array([x[0], x[1]])

    kf = kalman_filter(lambda t0,t: (tr, Q_lo), lambda t: (obs, None))
    st = kf.init(mu, cov)
    step = jax.jit(lambda s: kf.step(s, observed=jnp.array([1.0, 1.0])))
    _ = step(st)  # warmup

    t0 = time.perf_counter()
    for _ in range(5):
        st_new, _ = step(st)
        st_new.mean.block_until_ready()
    t_lo = (time.perf_counter() - t0) / 5 * 1000

    # Dense
    tr_d = tr.as_array()
    Q_d = Q_lo.as_array()
    obs_d = obs.as_array()
    kf_d = kalman_filter(lambda t0,t: (tr_d, Q_d), lambda t: (obs_d, None))
    st_d = kf_d.init(mu, cov)
    step_d = jax.jit(lambda s: kf_d.step(s, observed=jnp.array([1.0, 1.0])))
    _ = step_d(st_d)

    t0 = time.perf_counter()
    for _ in range(5):
        st_d_new, _ = step_d(st_d)
        st_d_new.mean.block_until_ready()
    t_dense = (time.perf_counter() - t0) / 5 * 1000

    mem = dim * dim * 8 / 1024**2
    marker = " <-- LO wins" if t_lo < t_dense else ""
    print(f"{dim:>6} | {t_dense:>9.1f}  | {t_lo:>9.1f}  | {mem:>7.0f}  {marker}")

   dim | dense (ms) |    LO (ms) | mem (MB)
---------------------------------------------


   100 |       0.4  |       0.2  |       0   <-- LO wins


   500 |       4.5  |       3.7  |       2   <-- LO wins


  1000 |      30.7  |      22.5  |       8   <-- LO wins


  2000 |     258.0  |     174.1  |      31   <-- LO wins


  3000 |     935.3  |     588.7  |      69   <-- LO wins

Summary¶

LinearOperators wrap v -> A @ v instead of storing the full matrix.
They compose via @, +, .T — same algebra as dense matrices.
Call .as_array() to materialize when you need the dense form.
The kalman_filter accepts both dense arrays and LinearOperators.
For small d (< 1000): dense is faster (direct LAPACK).
For large d (> 1000): LinearOperators save memory and use matrix-free solvers (CG, Lanczos) internally.