optax integration

EKI, EKS, and UKI are exposed as optax.GradientTransformation pairs so you can mix derivative-free ensemble updates with the rest of the optax ecosystem (schedules, gradient clipping, hybrid pipelines with gradient-based optimisers, …). The three constructors live in filterax.optax.

Constructors

Ensemble Kalman processes exposed as optax gradient transformations.

Each constructor returns an optax.GradientTransformation with the ensemble (or sigma-point belief) living inside the optax OptState and the running mean estimate flowing back through apply_updates.

• :func:eki — Ensemble Kalman Inversion.
• :func:eks — Ensemble Kalman Sampler (posterior sampling).
• :func:uki — Unscented Kalman Inversion (parametric).

Compose with optax.chain for clipping, schedules, hybrid pipelines with gradient-based optimisers, etc. See docs/design_docs/features/optax_ekp.md for the contract and examples.

eki(forward_fn: ForwardFn, obs: Float[Array, ' N_d'], noise_cov: lx.AbstractLinearOperator, *, n_ensemble: int = 50, init_spread: float = 1.0, scheduler: AbstractScheduler | None = None, key: PRNGKeyArray | int = 0) -> optax.GradientTransformation

Ensemble Kalman Inversion as an optax.GradientTransformation.

The grads argument to :meth:update is ignored — the ensemble provides search directions via the cross-covariance. The increment returned is the change in the ensemble mean, unflattened back to the params tree shape so :func:optax.apply_updates works as usual.

Parameters:

Name	Type	Description	Default
forward_fn	ForwardFn	θ → G(θ) mapping. Closure-captured (not stored in state) so it doesn't have to be a pytree leaf.	required
obs	Float[Array, ' N_d']	Observation vector y ∈ ℝ^{N_d}.	required
noise_cov	AbstractLinearOperator	Observation noise covariance Γ.	required
n_ensemble	int	Ensemble size J.	50
init_spread	float	Standard deviation of the initial ensemble perturbation around params (the seed mean).	1.0
scheduler	AbstractScheduler | None	Step-size strategy. Defaults to :class:~filterax.DataMisfitController.	None
key	PRNGKeyArray | int	PRNG seed (int or PRNG array) for the initial ensemble.	0

Returns:

Type	Description
GradientTransformation	optax.GradientTransformation whose update advances the
GradientTransformation	ensemble and returns the per-step mean delta as the optax
GradientTransformation	update.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import optax
>>> from filterax import FixedScheduler
>>> from filterax.optax import eki
>>> G = jnp.array([[1.0, 0.5], [-0.3, 1.2]])
>>> transform = eki(
...     forward_fn=lambda theta: G @ theta,
...     obs=jnp.array([0.5, 0.9]),
...     noise_cov=lx.DiagonalLinearOperator(0.1 * jnp.ones(2)),
...     n_ensemble=8,
...     scheduler=FixedScheduler(dt=1.0),
... )
>>> params = jnp.zeros(2)
>>> state = transform.init(params)
>>> updates, state = transform.update(None, state, params)
>>> params = optax.apply_updates(params, updates)
>>> params.shape
(2,)

Source code in src/filterax/_processes/_optax.py

def eki(
    forward_fn: ForwardFn,
    obs: Float[Array, " N_d"],
    noise_cov: lx.AbstractLinearOperator,
    *,
    n_ensemble: int = 50,
    init_spread: float = 1.0,
    scheduler: AbstractScheduler | None = None,
    key: PRNGKeyArray | int = 0,
) -> optax.GradientTransformation:
    r"""Ensemble Kalman Inversion as an ``optax.GradientTransformation``.

    The ``grads`` argument to :meth:`update` is ignored — the ensemble
    provides search directions via the cross-covariance. The increment
    returned is the change in the *ensemble mean*, unflattened back to
    the ``params`` tree shape so :func:`optax.apply_updates` works as
    usual.

    Args:
        forward_fn: ``θ → G(θ)`` mapping. Closure-captured (not stored
            in state) so it doesn't have to be a pytree leaf.
        obs: Observation vector ``y ∈ ℝ^{N_d}``.
        noise_cov: Observation noise covariance ``Γ``.
        n_ensemble: Ensemble size ``J``.
        init_spread: Standard deviation of the initial ensemble
            perturbation around ``params`` (the seed mean).
        scheduler: Step-size strategy. Defaults to
            :class:`~filterax.DataMisfitController`.
        key: PRNG seed (int or PRNG array) for the initial ensemble.

    Returns:
        ``optax.GradientTransformation`` whose ``update`` advances the
        ensemble and returns the per-step mean delta as the optax
        update.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import optax
        >>> from filterax import FixedScheduler
        >>> from filterax.optax import eki
        >>> G = jnp.array([[1.0, 0.5], [-0.3, 1.2]])
        >>> transform = eki(
        ...     forward_fn=lambda theta: G @ theta,
        ...     obs=jnp.array([0.5, 0.9]),
        ...     noise_cov=lx.DiagonalLinearOperator(0.1 * jnp.ones(2)),
        ...     n_ensemble=8,
        ...     scheduler=FixedScheduler(dt=1.0),
        ... )
        >>> params = jnp.zeros(2)
        >>> state = transform.init(params)
        >>> updates, state = transform.update(None, state, params)
        >>> params = optax.apply_updates(params, updates)
        >>> params.shape
        (2,)
    """
    sched = scheduler if scheduler is not None else DataMisfitController()
    base_key = jr.PRNGKey(key) if isinstance(key, int) else key

    def init_fn(params: Any) -> EKIOptaxState:
        mean, unravel = _flatten(params)
        # θ⁽ʲ⁾ ~ 𝒩(mean, init_spread² I), recentred so that
        # ``particles.mean(axis=0) == mean`` exactly — otherwise the
        # first ``update`` would return a non-zero delta driven purely
        # by finite-sample sampling noise rather than by the data.
        noise = init_spread * jr.normal(
            base_key, (n_ensemble, mean.shape[0]), dtype=mean.dtype
        )
        noise = noise - jnp.mean(noise, axis=0, keepdims=True)
        particles = mean[None, :] + noise
        return EKIOptaxState(
            particles=particles,
            algo_time=jnp.asarray(0.0, dtype=mean.dtype),
            step=jnp.asarray(0, dtype=jnp.int32),
            unravel=unravel,
        )

    def update_fn(
        grads: Any,
        state: EKIOptaxState,
        params: Any | None = None,
    ) -> tuple[Any, EKIOptaxState]:
        del grads
        if params is None:
            raise ValueError(
                "filterax.optax.eki requires `params` to compute the delta."
            )
        prev_mean, _ = _flatten(params)

        evals = jax.vmap(forward_fn)(state.particles)
        shim = _scheduler_state_shim(
            state.particles, evals, obs, noise_cov, state.step, state.algo_time
        )
        dt = sched.get_dt(shim)
        delta = _eki_delta(state.particles, evals, obs, noise_cov, dt)
        particles_new = state.particles + delta
        new_mean = jnp.mean(particles_new, axis=0)
        # Unflatten the mean increment back into the params tree.
        mean_update = state.unravel(new_mean - prev_mean)
        return mean_update, EKIOptaxState(
            particles=particles_new,
            algo_time=state.algo_time + dt,
            step=state.step + 1,
            unravel=state.unravel,
        )

    # optax's TransformInitFn/TransformUpdateFn protocols assume the
    # state is array-typed; our richer eqx.Module state is functionally
    # compatible (pytree + JIT-stable) but the narrow stub rejects it.
    return optax.GradientTransformation(init_fn, update_fn)  # ty: ignore[invalid-argument-type]

eks(forward_fn: ForwardFn, obs: Float[Array, ' N_d'], noise_cov: lx.AbstractLinearOperator, *, n_ensemble: int = 100, init_spread: float = 1.0, scheduler: AbstractScheduler | None = None, key: PRNGKeyArray | int = 0) -> optax.GradientTransformation

Ensemble Kalman Sampler as an optax.GradientTransformation.

Adds Brownian noise so the ensemble samples the posterior rather than collapsing. Access state.particles after burn-in for posterior samples. The PRNG key advances at every step via :func:jr.fold_in so successive draws are independent.

Parameters:

Name	Type	Description	Default
forward_fn	ForwardFn	θ → G(θ).	required
obs	Float[Array, ' N_d']	Observation vector.	required
noise_cov	AbstractLinearOperator	Observation noise covariance Γ.	required
n_ensemble	int	Ensemble size J (recommended ≥ 100 for sampling).	100
init_spread	float	Standard deviation of the initial ensemble perturbation around params.	1.0
scheduler	AbstractScheduler | None	Step-size strategy. Use :class:~filterax.EKSStableScheduler for stability.	None
key	PRNGKeyArray | int	PRNG seed.	0

Source code in src/filterax/_processes/_optax.py

def eks(
    forward_fn: ForwardFn,
    obs: Float[Array, " N_d"],
    noise_cov: lx.AbstractLinearOperator,
    *,
    n_ensemble: int = 100,
    init_spread: float = 1.0,
    scheduler: AbstractScheduler | None = None,
    key: PRNGKeyArray | int = 0,
) -> optax.GradientTransformation:
    r"""Ensemble Kalman Sampler as an ``optax.GradientTransformation``.

    Adds Brownian noise so the ensemble samples the posterior rather
    than collapsing. Access ``state.particles`` after burn-in for
    posterior samples. The PRNG key advances at every step via
    :func:`jr.fold_in` so successive draws are independent.

    Args:
        forward_fn: ``θ → G(θ)``.
        obs: Observation vector.
        noise_cov: Observation noise covariance ``Γ``.
        n_ensemble: Ensemble size ``J`` (recommended ≥ 100 for sampling).
        init_spread: Standard deviation of the initial ensemble
            perturbation around ``params``.
        scheduler: Step-size strategy. Use
            :class:`~filterax.EKSStableScheduler` for stability.
        key: PRNG seed.
    """
    sched = scheduler if scheduler is not None else EKSStableScheduler()
    base_key = jr.PRNGKey(key) if isinstance(key, int) else key
    init_key, noise_key = jr.split(base_key)

    def init_fn(params: Any) -> EKSOptaxState:
        mean, unravel = _flatten(params)
        noise = init_spread * jr.normal(
            init_key, (n_ensemble, mean.shape[0]), dtype=mean.dtype
        )
        particles = mean[None, :] + noise
        return EKSOptaxState(
            particles=particles,
            algo_time=jnp.asarray(0.0, dtype=mean.dtype),
            step=jnp.asarray(0, dtype=jnp.int32),
            key=noise_key,
            unravel=unravel,
        )

    def update_fn(
        grads: Any,
        state: EKSOptaxState,
        params: Any | None = None,
    ) -> tuple[Any, EKSOptaxState]:
        del grads
        if params is None:
            raise ValueError(
                "filterax.optax.eks requires `params` to compute the delta."
            )
        prev_mean, _ = _flatten(params)

        J, N_p = state.particles.shape
        evals = jax.vmap(forward_fn)(state.particles)
        shim = _scheduler_state_shim(
            state.particles, evals, obs, noise_cov, state.step, state.algo_time
        )
        dt = sched.get_dt(shim)

        # EKI-style drift + mean-pull + ensemble-preconditioned noise.
        drift = _eki_delta(state.particles, evals, obs, noise_cov, dt)
        mean = jnp.mean(state.particles, axis=0)
        anom = state.particles - mean[None, :]
        drift_pull = -((N_p + 1) / J) * dt * anom
        step_key = jr.fold_in(state.key, state.step)
        xi = jr.normal(step_key, (J, J), dtype=state.particles.dtype)
        brownian = (
            jnp.sqrt(2.0 * dt) * einx.dot("j k, k p -> j p", xi, anom) / jnp.sqrt(J - 1)
        )

        particles_new = state.particles + drift + drift_pull + brownian
        new_mean = jnp.mean(particles_new, axis=0)
        mean_update = state.unravel(new_mean - prev_mean)
        return mean_update, EKSOptaxState(
            particles=particles_new,
            algo_time=state.algo_time + dt,
            step=state.step + 1,
            key=state.key,
            unravel=state.unravel,
        )

    # optax's TransformInitFn/TransformUpdateFn protocols assume the
    # state is array-typed; our richer eqx.Module state is functionally
    # compatible (pytree + JIT-stable) but the narrow stub rejects it.
    return optax.GradientTransformation(init_fn, update_fn)  # ty: ignore[invalid-argument-type]

uki(forward_fn: ForwardFn, obs: Float[Array, ' N_d'], noise_cov: lx.AbstractLinearOperator, *, init_cov: float | Float[Array, 'N_p N_p'] = 1.0, scheduler: AbstractScheduler | None = None, alpha: float = 1.0, beta: float = 2.0, kappa: float = 0.0) -> optax.GradientTransformation

Unscented Kalman Inversion as an optax.GradientTransformation.

Parametric: state carries (μ, Σ) plus the unflatten callback. Each update regenerates 2 Nₚ + 1 sigma points, evaluates forward_fn on them, and applies the unscented Kalman update. Access state.covariance for the posterior covariance estimate.

Parameters:

Name	Type	Description	Default
forward_fn	ForwardFn	θ → G(θ).	required
obs	Float[Array, ' N_d']	Observation vector.	required
noise_cov	AbstractLinearOperator	Observation noise covariance Γ.	required
init_cov	float | Float[Array, 'N_p N_p']	Initial parameter covariance. Scalar → σ² I; matrix → used directly.	1.0
scheduler	AbstractScheduler | None	Step-size strategy.	None
alpha	float	Sigma-point spread parameter (Wan & van der Merwe 2000).	1.0
beta	float	Prior-moment weighting (2.0 is optimal for Gaussian).	2.0
kappa	float	Secondary scaling parameter; typically 0 or 3 − Nₚ.	0.0

Source code in src/filterax/_processes/_optax.py

def uki(
    forward_fn: ForwardFn,
    obs: Float[Array, " N_d"],
    noise_cov: lx.AbstractLinearOperator,
    *,
    init_cov: float | Float[Array, "N_p N_p"] = 1.0,
    scheduler: AbstractScheduler | None = None,
    alpha: float = 1.0,
    beta: float = 2.0,
    kappa: float = 0.0,
) -> optax.GradientTransformation:
    r"""Unscented Kalman Inversion as an ``optax.GradientTransformation``.

    Parametric: state carries ``(μ, Σ)`` plus the unflatten callback.
    Each ``update`` regenerates ``2 Nₚ + 1`` sigma points, evaluates
    ``forward_fn`` on them, and applies the unscented Kalman update.
    Access ``state.covariance`` for the posterior covariance estimate.

    Args:
        forward_fn: ``θ → G(θ)``.
        obs: Observation vector.
        noise_cov: Observation noise covariance ``Γ``.
        init_cov: Initial parameter covariance. Scalar → ``σ² I``;
            matrix → used directly.
        scheduler: Step-size strategy.
        alpha: Sigma-point spread parameter (Wan & van der Merwe 2000).
        beta: Prior-moment weighting (``2.0`` is optimal for Gaussian).
        kappa: Secondary scaling parameter; typically ``0`` or ``3 − Nₚ``.
    """
    sched = scheduler if scheduler is not None else DataMisfitController()

    def init_fn(params: Any) -> UKIOptaxState:
        mean, unravel = _flatten(params)
        if jnp.ndim(jnp.asarray(init_cov)) == 0:
            Sigma = (init_cov**2) * jnp.eye(mean.shape[0], dtype=mean.dtype)
        else:
            Sigma = jnp.asarray(init_cov)
        cov_op = lx.MatrixLinearOperator(Sigma, tags=lx.positive_semidefinite_tag)
        return UKIOptaxState(
            mean=mean,
            covariance=cov_op,
            algo_time=jnp.asarray(0.0, dtype=mean.dtype),
            step=jnp.asarray(0, dtype=jnp.int32),
            unravel=unravel,
        )

    def update_fn(
        grads: Any,
        state: UKIOptaxState,
        params: Any | None = None,
    ) -> tuple[Any, UKIOptaxState]:
        del grads
        if params is None:
            raise ValueError(
                "filterax.optax.uki requires `params` to compute the delta."
            )
        prev_mean, _ = _flatten(params)

        points, w_mean, w_cov = sigma_points(
            state.mean,
            state.covariance,
            alpha=alpha,
            beta=beta,
            kappa=kappa,
        )
        evals = jax.vmap(forward_fn)(points)
        shim = _scheduler_state_shim(
            points, evals, obs, noise_cov, state.step, state.algo_time
        )
        dt = sched.get_dt(shim)

        # Unscented Kalman update.
        y_mean = einx.dot("k, k d -> d", w_mean, evals)
        theta_anom = einx.subtract("k p, p -> k p", points, state.mean)
        y_anom = einx.subtract("k d, d -> k d", evals, y_mean)
        weighted_theta = einx.multiply("k p, k -> k p", theta_anom, w_cov)
        C_theta_y = einx.dot("k p, k d -> p d", weighted_theta, y_anom)
        weighted_y = einx.multiply("k d, k -> k d", y_anom, w_cov)
        S_dense = einx.dot("k a, k b -> a b", weighted_y, y_anom)
        S_tempered = S_dense + _safe_inv_dt(dt) * noise_cov.as_matrix()
        innovation = obs - y_mean
        K_gain = jnp.linalg.solve(S_tempered.T, C_theta_y.T).T

        new_mean = state.mean + einx.dot("p d, d -> p", K_gain, innovation)
        # Σ_{n+1} = Σ̂ − K S_tempered Kᵀ; Δt enters once, through Δt⁻¹ Γ in S.
        Sigma_dense = state.covariance.as_matrix()
        K_S = einx.dot("p a, a b -> p b", K_gain, S_tempered)
        Sigma_new = Sigma_dense - einx.dot("p d, q d -> p q", K_S, K_gain)
        Sigma_new = 0.5 * (Sigma_new + Sigma_new.T)
        cov_new = lx.MatrixLinearOperator(Sigma_new, tags=lx.positive_semidefinite_tag)

        mean_update = state.unravel(new_mean - prev_mean)
        return mean_update, UKIOptaxState(
            mean=new_mean,
            covariance=cov_new,
            algo_time=state.algo_time + dt,
            step=state.step + 1,
            unravel=state.unravel,
        )

    # optax's TransformInitFn/TransformUpdateFn protocols assume the
    # state is array-typed; our richer eqx.Module state is functionally
    # compatible (pytree + JIT-stable) but the narrow stub rejects it.
    return optax.GradientTransformation(init_fn, update_fn)  # ty: ignore[invalid-argument-type]

Contract

The grads argument to update(grads, state, params) is ignored — the ensemble (or sigma-point quadrature) provides search directions. The returned updates is the change in the ensemble mean, unflattened back into the original params PyTree shape so optax.apply_updates works without modification.

The ensemble itself lives in the optax OptState; params only ever holds the running mean. Pytree-shaped params (nested dicts, NamedTuples, …) round-trip via jax.flatten_util.ravel_pytree.

Patterns

Basic EKI loop

import filterax as flx
import optax

transform = flx.optax.eki(
    forward_fn=simulator,
    obs=observations,
    noise_cov=Gamma,
    scheduler=flx.DataMisfitController(),
)

params = initial_params
state = transform.init(params)
for _ in range(50):
    updates, state = transform.update(None, state, params)
    params = optax.apply_updates(params, updates)

Posterior sampling with EKS

sampler = flx.optax.eks(
    forward_fn=simulator,
    obs=observations,
    noise_cov=Gamma,
    n_ensemble=200,
)
state = sampler.init(params)
samples = []
for step in range(500):
    updates, state = sampler.update(None, state, params)
    params = optax.apply_updates(params, updates)
    if step > burnin:
        samples.append(state.particles)  # (200, Nₚ) per step

Hybrid: gradient warm-start + EKI refinement

# Phase 1 — Adam handles fast initial descent with gradients.
adam = optax.adam(1e-3)
state_a = adam.init(params)
for _ in range(1000):
    grads = jax.grad(loss_fn)(params)
    updates, state_a = adam.update(grads, state_a, params)
    params = optax.apply_updates(params, updates)

# Phase 2 — switch to EKI for gradient-free refinement.
phase2 = flx.optax.eki(
    forward_fn=simulator,
    obs=observations,
    noise_cov=Gamma,
    init_spread=0.1,  # tight cluster around the warm-started mean
)
state_b = phase2.init(params)
for _ in range(50):
    updates, state_b = phase2.update(None, state_b, params)
    params = optax.apply_updates(params, updates)

Composition with other optax transforms

optimizer = optax.chain(
    flx.optax.eki(forward_fn=sim, obs=y, noise_cov=R),
    optax.clip_by_global_norm(1.0),
)
state = optimizer.init(params)
updates, state = optimizer.update(None, state, params)
params = optax.apply_updates(params, updates)

Internal state types

These are returned from the optax transforms' init and threaded through update. Inspect them directly for the full ensemble (EKI / EKS) or parametric belief (UKI) — useful for diagnostics and posterior sampling.

Ensemble Kalman processes exposed as optax gradient transformations.

Each constructor returns an optax.GradientTransformation with the ensemble (or sigma-point belief) living inside the optax OptState and the running mean estimate flowing back through apply_updates.

• :func:eki — Ensemble Kalman Inversion.
• :func:eks — Ensemble Kalman Sampler (posterior sampling).
• :func:uki — Unscented Kalman Inversion (parametric).

Compose with optax.chain for clipping, schedules, hybrid pipelines with gradient-based optimisers, etc. See docs/design_docs/features/optax_ekp.md for the contract and examples.

EKIOptaxState

Bases: Module

Internal optax state for :func:eki.

Source code in src/filterax/_processes/_optax.py

class EKIOptaxState(eqx.Module, strict=True):
    """Internal optax state for :func:`eki`."""

    particles: Float[Array, "J N_p"]
    algo_time: Float[Array, ""]
    step: jnp.ndarray
    unravel: Callable[[Float[Array, " N_p"]], Any] = eqx.field(static=True)

EKSOptaxState

Bases: Module

Internal optax state for :func:eks.

Source code in src/filterax/_processes/_optax.py

class EKSOptaxState(eqx.Module, strict=True):
    """Internal optax state for :func:`eks`."""

    particles: Float[Array, "J N_p"]
    algo_time: Float[Array, ""]
    step: jnp.ndarray
    key: PRNGKeyArray
    unravel: Callable[[Float[Array, " N_p"]], Any] = eqx.field(static=True)

UKIOptaxState

Bases: Module

Internal optax state for :func:uki.

Source code in src/filterax/_processes/_optax.py

class UKIOptaxState(eqx.Module, strict=True):
    """Internal optax state for :func:`uki`."""

    mean: Float[Array, " N_p"]
    covariance: lx.AbstractLinearOperator
    algo_time: Float[Array, ""]
    step: jnp.ndarray
    unravel: Callable[[Float[Array, " N_p"]], Any] = eqx.field(static=True)