"""JAX backend: wraps the CIF obs operator via ``jax.custom_vjp``.
The forward and backward passes call the numpy CIF operator directly.
This works in JAX eager mode; for JIT compilation the callbacks would need
to be wrapped in ``jax.pure_callback`` (requires knowing output shapes ahead
of time — see the ``_jit_wrap`` comment below).
Training uses ``optax`` for optimisation and assumes a Flax-style model API:
``params = model.init(key, x)``
``y = model.apply(params, x)``
"""
from __future__ import annotations
import numpy as np
from typing import Any, Callable, List, Optional, Tuple
from .base import SurrogateBackend
[docs]
class JaxBackend(SurrogateBackend):
"""JAX surrogate backend."""
def __init__(self, obs_operator: Any) -> None:
super().__init__(obs_operator)
# Cached output shape for pure_callback (filled on first wrap_operator call)
self._out_shape: Optional[Tuple[int, ...]] = None
# ------------------------------------------------------------------
[docs]
def build_default_model(
self,
n_state: int,
n_obs: int,
hidden_size: int,
n_hidden_layers: int,
) -> Any:
"""Return a Flax ``linen.Module`` MLP: n_state -> hidden -> n_obs."""
try:
import flax.linen as nn
except ImportError as exc:
raise ImportError(
"Flax is required for the jax backend: pip install flax"
) from exc
features = [hidden_size] * n_hidden_layers + [n_obs]
class _MLP(nn.Module):
@nn.compact
def __call__(self, x):
for feat in features[:-1]:
x = nn.relu(nn.Dense(feat)(x))
return nn.Dense(features[-1])(x)
return _MLP()
# ------------------------------------------------------------------
[docs]
def wrap_operator(self, run_mode: str = "fwd",
reload_results: bool = False) -> Callable:
"""Return a JAX-differentiable callable wrapping the CIF operator.
Args:
run_mode: ``"fwd"`` uses ``obs_operator.forward``; ``"tl"`` uses
``obs_operator.tangent_linear``. The backward pass always
calls ``obs_operator.adjoint``.
reload_results: Forwarded to every operator call.
Note: not JIT-compilable as-is. For JIT support replace the direct
numpy calls with ``jax.pure_callback``; see the inline comment.
"""
try:
import jax
import jax.numpy as jnp
except ImportError as exc:
raise ImportError(
"JAX is required for the jax backend: pip install jax"
) from exc
obs_operator = self.obs_operator
fwd_call = (obs_operator.tangent_linear
if run_mode == "tl" else obs_operator.forward)
# ------------------------------------------------------------------
# JIT note: to support jax.jit, replace the body of H and H_bwd with
# pure_callback calls, e.g.:
#
# out_shape = jax.ShapeDtypeStruct(y0.shape, jnp.float64)
# return jax.pure_callback(
# lambda x_cb: fwd_call(np.asarray(x_cb)).astype(np.float64),
# out_shape, x,
# )
#
# This requires knowing the output shape before tracing, which can be
# determined by a one-time probe call stored in self._out_shape.
# ------------------------------------------------------------------
@jax.custom_vjp
def H(x: "jnp.ndarray") -> "jnp.ndarray":
y_np = fwd_call(np.asarray(x, dtype=np.float64),
reload_results=reload_results)
return jnp.asarray(y_np, dtype=x.dtype)
def H_fwd(x: "jnp.ndarray") -> Tuple["jnp.ndarray", "jnp.ndarray"]:
return H(x), x # residuals: save x for backward
def H_bwd(
x_saved: "jnp.ndarray", g: "jnp.ndarray"
) -> Tuple["jnp.ndarray"]:
dx_np = obs_operator.adjoint(
np.asarray(x_saved, dtype=np.float64),
np.asarray(g, dtype=np.float64),
reload_results=reload_results,
)
return (jnp.asarray(dx_np, dtype=x_saved.dtype),)
H.defvjp(H_fwd, H_bwd)
return H
# ------------------------------------------------------------------
[docs]
def train(
self,
model: Any,
sample_fn: Callable[[int], Any],
n_steps: int = 5000,
*,
run_mode: str = "fwd",
reload_results: bool = False,
adjoint_consistency: bool = False,
lambda_adj: float = 0.1,
) -> Tuple[Any, List[float]]:
"""Train *model* (Flax-style) against the live CIF operator.
Args:
model: A Flax ``linen.Module`` (or any module with ``.init`` /
``.apply`` methods).
sample_fn: ``callable(batch_size) -> jnp.ndarray`` returning a
single state vector of shape ``(N_state,)``.
n_steps: Number of gradient steps.
run_mode: Forwarded to :meth:`wrap_operator`; ``"tl"`` trains
against the tangent-linear operator instead of the full forward.
reload_results: Forwarded to :meth:`wrap_operator`.
adjoint_consistency: Add VJP-agreement loss term (Mode 3).
lambda_adj: Weight of the adjoint consistency loss.
Returns:
``(trained_params, losses)``
"""
try:
import jax
import jax.numpy as jnp
import optax
except ImportError as exc:
raise ImportError(
"JAX and optax are required for JAX training: "
"pip install jax optax"
) from exc
H = self.wrap_operator(run_mode=run_mode, reload_results=reload_results)
obs_operator = self.obs_operator
dummy = jnp.asarray(sample_fn(1), dtype=jnp.float32)
params = model.init(jax.random.PRNGKey(0), dummy)
optimizer = optax.adam(1e-3)
opt_state = optimizer.init(params)
def loss_fn(params: Any, x: "jnp.ndarray") -> "jnp.ndarray":
y_target = H(x)
y_pred = model.apply(params, x)
fwd_loss = jnp.mean((y_pred - y_target) ** 2)
if not adjoint_consistency:
return fwd_loss
# Mode 3: VJP agreement — match surrogate's J^T dy to CIF adjoint
key = jax.random.PRNGKey(0)
dy = jax.random.normal(key, y_target.shape)
# Surrogate VJP w.r.t. x
_, vjp_x = jax.vjp(lambda xi: model.apply(params, xi), x)
(dx_surr,) = vjp_x(dy)
# CIF adjoint (numpy, opaque to JAX)
dx_cif = jnp.asarray(
obs_operator.adjoint(
np.asarray(x, dtype=np.float64),
np.asarray(dy, dtype=np.float64),
),
dtype=x.dtype,
)
adj_loss = jnp.mean((dx_surr - dx_cif) ** 2)
return fwd_loss + lambda_adj * adj_loss
@jax.jit
def step(
params: Any, opt_state: Any, x: "jnp.ndarray"
) -> Tuple[Any, Any, "jnp.ndarray"]:
loss_val, grads = jax.value_and_grad(loss_fn)(params, x)
updates, new_opt_state = optimizer.update(grads, opt_state)
new_params = optax.apply_updates(params, updates)
return new_params, new_opt_state, loss_val
losses: List[float] = []
for _ in range(n_steps):
x = jnp.asarray(sample_fn(1), dtype=jnp.float32)
params, opt_state, loss_val = step(params, opt_state, x)
losses.append(float(loss_val))
return params, losses
# ------------------------------------------------------------------
[docs]
def export_onnx(
self, path: str, model: Any, input_shape: Tuple[int, ...]
) -> None:
"""Export to ONNX is not directly supported from JAX.
Recommended conversion path::
# Step 1 – JAX -> TensorFlow SavedModel
import jax
from jax.experimental import jax2tf
import tensorflow as tf
params = ... # trained params from train()
def predict(x):
return model.apply(params, x)
tf_fn = tf.function(
jax2tf.convert(predict, enable_xla=False),
input_signature=[
tf.TensorSpec(shape=input_shape, dtype=tf.float32)
],
)
tf.saved_model.save(tf_fn, "/tmp/jax_saved_model")
# Step 2 – SavedModel -> ONNX (shell)
# python -m tf2onnx.convert \\
# --saved-model /tmp/jax_saved_model \\
# --output surrogate.onnx
References:
https://github.com/google/jax/tree/main/jax/experimental#jax2tf-converter
https://github.com/onnx/tensorflow-onnx
"""
raise NotImplementedError(
"JAX -> ONNX export requires the jax2tf + tf2onnx pipeline. "
"See the docstring of JaxBackend.export_onnx for the conversion "
"skeleton, or use TorchBackend for direct ONNX export."
)
