From 9bb06160d0871c54390c170e6cce8054ebeffca6 Mon Sep 17 00:00:00 2001 From: Alexander Fengler Date: Wed, 5 Aug 2026 20:10:52 -0400 Subject: [PATCH] feat: jax -> ONNX export via jax2onnx; jaxtrain now produces HSSM artifacts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit jaxtrain was a dead end for the ecosystem: it saved flax parameter bytes that nothing downstream could read — transform-onnx only handles torch state dicts and HSSM consumes ONNX exclusively. This closes the single missing edge. - New lanfactory/onnx/jax_export.py: export_forward_to_onnx (shared core) + transform_jax_to_onnx (file-based, for retroactive conversion of existing .jax artifacts), exposed as the `transform-jax-onnx` CLI - The jax trainer exports ONNX alongside its other artifacts (mirroring the torch trainer's _save_onnx) and logs it to MLflow; `--no-export-onnx` on jaxtrain opts out - Dependency: jax2onnx==0.15.* (exact-pinned 0.x line), opset 17 to match the sbi/bayesflow exporters Two contract decisions, both verified empirically: - The graph is traced with a concrete (1, input_dim) dummy — every dim static, zero dynamic axes — matching the torch MLP exporter and the production networks on franklab/HSSM, NOT rank-1: jax2onnx lowers nn.Dense to Gemm, whose ONNX spec requires rank-2 inputs (a rank-1 trace is rejected by onnxruntime). Verified end-to-end that HSSM's make_jax_func loads the export and rank-1-per-trial + jax.vmap consumption matches exactly. - The exported graph is the EVAL head: identity for logprob (LAN), logsigmoid for logits (CPN/OPN) — matching every torch export path (_save_onnx calls .eval(); transform-onnx exports under torch's EVAL default) and HSSM's element-wise log-likelihood consumption. Exporting the raw training head would silently corrupt every CPN/OPN logp by +log(1+exp(-logit)); a test asserts the logits export equals the logsigmoid head and is NOT the raw head. Tests: contract (all dims concrete, (1,D) shape, op allowlist calibrated against jax2onnx 0.15's actual lowering — relu emits Max, the logsigmoid head emits Neg/Exp/Add/Log), parity (ORT vs eval-head jax forward over 1000 draws, atol 1e-4), file-based vs in-trainer export equivalence, logits-head correctness, and the opt-out flag. Co-Authored-By: Claude Opus 5 --- pyproject.toml | 5 + src/lanfactory/cli/jax_train.py | 7 + src/lanfactory/onnx/__init__.py | 2 + src/lanfactory/onnx/jax_export.py | 150 ++++++++++++++ src/lanfactory/trainers/jax_mlp.py | 39 ++++ tests/onnx/__init__.py | 0 tests/onnx/test_jax_export.py | 320 +++++++++++++++++++++++++++++ 7 files changed, 523 insertions(+) create mode 100644 src/lanfactory/onnx/jax_export.py create mode 100644 tests/onnx/__init__.py create mode 100644 tests/onnx/test_jax_export.py diff --git a/pyproject.toml b/pyproject.toml index ebec8f3..a46caf4 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,10 @@ dependencies = [ "tqdm>=4.67.1", "frozendict>=2.4.6", "onnx>=1.17.0", + # jax -> ONNX export for jaxtrain networks. Exact-pinned to the 0.15 line: + # 0.x releases move the API, and the parity test in tests/onnx guards the + # ecosystem's single-trial ONNX contract against silent changes. + "jax2onnx==0.15.*", "matplotlib>=3.10.1", "plotly>=6.0.0", "seaborn>=0.13.2", @@ -156,6 +160,7 @@ exclude_lines = [ jaxtrain = "lanfactory.cli.jax_train:app" torchtrain = "lanfactory.cli.torch_train:app" transform-onnx = "lanfactory.onnx.transform_onnx:app" +transform-jax-onnx = "lanfactory.onnx.jax_export:app" upload-hf = "lanfactory.cli.upload_hf:app" download-hf = "lanfactory.cli.download_hf:app" diff --git a/src/lanfactory/cli/jax_train.py b/src/lanfactory/cli/jax_train.py index 213b5ca..93069f5 100644 --- a/src/lanfactory/cli/jax_train.py +++ b/src/lanfactory/cli/jax_train.py @@ -53,6 +53,12 @@ def main( help="Validate the pipeline without training. Useful for testing configurations.", is_flag=True, ), + export_onnx: bool = typer.Option( + True, + "--export-onnx/--no-export-onnx", + help="Export the trained network to ONNX (single-trial contract) " + "alongside the flax state. ONNX is the artifact HSSM consumes.", + ), mlflow_run_name: str = typer.Option( None, "--mlflow-run-name", @@ -525,6 +531,7 @@ def main( # Pass explicitly: inference from train_output_type mislabels OPN as # cpn (both train on logits). network_type=network_config["network_type"], + export_onnx=export_onnx, ) # ------------------------------------------------------------- diff --git a/src/lanfactory/onnx/__init__.py b/src/lanfactory/onnx/__init__.py index 0b25cb2..626293d 100755 --- a/src/lanfactory/onnx/__init__.py +++ b/src/lanfactory/onnx/__init__.py @@ -1,9 +1,11 @@ from .bayesflow import transform_bayesflow_to_onnx +from .jax_export import transform_jax_to_onnx from .sbi import transform_sbi_to_onnx from .transform_onnx import transform_to_onnx __all__ = [ "transform_to_onnx", + "transform_jax_to_onnx", "transform_sbi_to_onnx", "transform_bayesflow_to_onnx", ] diff --git a/src/lanfactory/onnx/jax_export.py b/src/lanfactory/onnx/jax_export.py new file mode 100644 index 0000000..2b4f397 --- /dev/null +++ b/src/lanfactory/onnx/jax_export.py @@ -0,0 +1,150 @@ +"""Export jaxtrain networks to ONNX. Can be run as a script. + +Closes the gap that made jaxtrain a dead end for HSSM: the jax trainer saves +flax parameter bytes (``*_train_state.jax``) which nothing downstream could +convert — ``transform-onnx`` only reads torch state dicts, and HSSM consumes +ONNX exclusively. + +The export follows the ecosystem's single-trial ONNX contract (see HSSMSpine +CLAUDE.md): the graph is traced with a **concrete** ``(1, input_dim)`` dummy +and no dynamic axes, exactly like the torch MLP exporter and the production +networks on franklab/HSSM. Every dim is static, so HSSM's load-time check +passes, and HSSM's rank-1-per-trial + ``jax.vmap`` consumption works because +the resulting graph is pure ``Gemm`` + elementwise activations (verified +end-to-end against ``hssm.make_jax_func``). Rank-1 ``(input_dim,)`` tracing is +NOT used here: jax2onnx lowers ``nn.Dense`` to ``Gemm``, whose ONNX spec +requires rank-2 inputs — a rank-1-traced graph is rejected by onnxruntime. +(The sbi/bayesflow exporters trace rank-1 because torch lowers Linear to +rank-agnostic MatMul+Add; different tracer, different constraint.) + +The exported graph is the **eval-mode** forward: for LANs (``logprob``) this +equals the raw head, and for CPN/OPN (``logits``) it applies the logsigmoid +transform ``-log(1 + exp(-x))`` so the network emits log choice probabilities. +This matches every torch export path (``_save_onnx`` calls ``.eval()``, and +``transform-onnx`` exports under ``torch.onnx``'s EVAL default) and is what +HSSM consumes as element-wise log-likelihood. Exporting the raw training head +for logits networks would silently corrupt every downstream logp by +``+log(1 + exp(-logit))``. +""" + +import pickle + +import typer + +# Opset matching the sbi/bayesflow exporters. The MLPs only need Gemm + +# elementwise activations, all ancient; a newer default (jax2onnx uses 23) +# would just narrow runtime compatibility (jaxonnxruntime) for no benefit. +DEFAULT_OPSET = 17 + + +def export_forward_to_onnx( + forward, + input_shape: int, + output_onnx_file: str, + model_name: str = "jax_mlp", + opset: int = DEFAULT_OPSET, +) -> None: + """Export a jax forward function to ONNX under the single-trial contract. + + The shared core for both the file-based CLI (``transform_jax_to_onnx``) + and the jax trainer's post-training export. ``forward`` must accept an + input of shape ``(1, input_shape)`` (flax Dense handles arbitrary leading + dims, so the trainers' forward functions qualify unchanged). + """ + import jax.numpy as jnp + from jax2onnx import to_onnx + + import onnx as onnx_lib + + model_proto = to_onnx( + forward, + # Concrete (1, D) dummy: every dim static — the load-bearing line of + # the ecosystem contract. See the module docstring for why (1, D) + # rather than rank-1 (Gemm requires rank 2). + inputs=[jnp.zeros((1, input_shape), dtype=jnp.float32)], + model_name=model_name, + opset=opset, + ) + onnx_lib.save_model(model_proto, output_onnx_file) + + +def transform_jax_to_onnx( + network_config_file: str, + state_file: str, + input_shape: int, + output_onnx_file: str, + opset: int = DEFAULT_OPSET, +) -> None: + """Transform a trained JaxMLP to ONNX format. + + Arguments + --------- + network_config_file (str): + Path to the pickle file containing the network configuration + (``layer_sizes``, ``activations``, ``train_output_type``). + state_file (str): + Path to the ``*_train_state.jax`` file written by the jax trainer + (flax ``to_bytes`` serialization of the parameters). + input_shape (int): + The size of the single-trial input vector for the model + (``n_params + 2`` for LANs). + output_onnx_file (str): + Path to the output ONNX file. + opset (int): + ONNX opset version to target. + """ + from lanfactory.trainers import JaxMLPFactory + + with open(network_config_file, "rb") as f: + network_config = pickle.load(f) + + # train=False: export the EVAL head. For logprob networks this equals the + # raw head; for logits networks it applies logsigmoid, matching the torch + # exporters and HSSM's log-likelihood consumption (see module docstring). + net = JaxMLPFactory(network_config=network_config, train=False) + forward, _ = net.make_forward_partial( + input_dim=input_shape, + state=state_file, + add_jitted=False, + ) + + export_forward_to_onnx( + forward, + input_shape=input_shape, + output_onnx_file=output_onnx_file, + model_name=str(network_config.get("network_type", "jax_mlp")), + opset=opset, + ) + + +app = typer.Typer() + + +def option_no_default(help: str) -> typer.Option: + return typer.Option(..., help=help, show_default=False) + + +@app.command() +def main( + network_config_file: str = option_no_default( + "Path to the network configuration file (pickle)." + ), + state_file: str = option_no_default( + "Path to the *_train_state.jax file (flax parameter bytes)." + ), + input_shape: int = option_no_default("Size of the input tensor for the model."), + output_onnx_file: str = option_no_default("Path to the output ONNX file."), + opset: int = typer.Option(DEFAULT_OPSET, help="ONNX opset version to target."), +): + """Convert a jaxtrain-produced JaxMLP to ONNX format.""" + transform_jax_to_onnx( + network_config_file, + state_file, + input_shape, + output_onnx_file, + opset=opset, + ) + + +if __name__ == "__main__": + app() diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py index 70befb4..beff289 100755 --- a/src/lanfactory/trainers/jax_mlp.py +++ b/src/lanfactory/trainers/jax_mlp.py @@ -501,6 +501,7 @@ def train_and_evaluate( save_outputs: bool = True, verbose: int = 1, network_type: str | None = None, + export_onnx: bool = True, ) -> train_state.TrainState: """Train and evaluate JAXMLP model. Arguments @@ -523,6 +524,9 @@ def train_and_evaluate( filenames. When None it is inferred from train_output_type — which cannot distinguish cpn from opn (both use logits), so callers that know the type should pass it. + export_onnx (bool): + Whether to export the trained network to ONNX alongside the + flax state (single-trial contract; the artifact HSSM consumes). Returns ------- flax.core.frozen_dict.FrozenDict: @@ -653,6 +657,39 @@ def train_and_evaluate( ) print("Saving training data details to: " + data_details_path) + # ONNX export (single-trial contract) — the artifact HSSM consumes. + # Mirrors the torch trainer's _save_onnx step; previously the jax + # path produced only flax bytes, which nothing downstream reads. + onnx_path = None + if export_onnx: + try: + from functools import partial as _partial + + from lanfactory.onnx.jax_export import export_forward_to_onnx + + # Export the EVAL head, not the raw training head: for + # logits networks (cpn/opn) eval applies logsigmoid, which + # is what the torch exporters emit and what HSSM consumes + # as log-likelihood. Parameters are head-independent, so + # the trained state applies to the eval twin unchanged. + eval_model = JaxMLP( + layer_sizes=self.model.layer_sizes, + activations=self.model.activations, + train_output_type=self.model.train_output_type, + train=False, + ) + onnx_path = f"{full_path}_model.onnx" + export_forward_to_onnx( + _partial(eval_model.apply, state.params), + input_shape=int(self.train_dl.dataset.input_dim), + output_onnx_file=onnx_path, + model_name=network_type, + ) + print("Saving ONNX export to: " + onnx_path) + except Exception as e: + onnx_path = None + print(f"Failed to export ONNX: {e}") + if self.mlflow_on: try: mlflow.log_artifact( @@ -665,6 +702,8 @@ def train_and_evaluate( mlflow.log_artifact( data_details_path, artifact_path="training_output" ) + if onnx_path is not None: + mlflow.log_artifact(onnx_path, artifact_path="training_output") mlflow.end_run() except Exception as e: print(f"Failed to log artifacts to MLflow: {e}") diff --git a/tests/onnx/__init__.py b/tests/onnx/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/onnx/test_jax_export.py b/tests/onnx/test_jax_export.py new file mode 100644 index 0000000..3ba6c0d --- /dev/null +++ b/tests/onnx/test_jax_export.py @@ -0,0 +1,320 @@ +"""Tests for the jax -> ONNX exporter (feat/jax-onnx-export). + +The parity test is the load-bearing one: it guards the ecosystem's +single-trial ONNX contract against silent changes in jax2onnx (0.x, pinned +but moving) and in our own wrapper. +""" + +import pickle +from pathlib import Path + +import lanfactory +import numpy as np +import pytest +from lanfactory.onnx import transform_jax_to_onnx +from torch.utils.data import DataLoader + +onnx = pytest.importorskip("onnx") +ort = pytest.importorskip("onnxruntime") + +INPUT_DIM = 6 + +# Calibrated against jax2onnx 0.15 lowering: relu -> Max, logsigmoid eval +# head -> Neg/Exp/Add/Log (+ Constant for literals). All implemented by +# jaxonnxruntime, HSSM's consumer. +ALLOWED_OPS = { + "Gemm", + "MatMul", + "Add", + "Tanh", + "Max", + "Sigmoid", + "Neg", + "Exp", + "Log", + "Constant", +} + +NETWORK_CONFIG = { + "layer_sizes": [16, 8, 1], + "activations": ["tanh", "tanh", "linear"], + "train_output_type": "logprob", + "network_type": "lan", +} + + +def make_training_pickle(path: Path, features_key="lan_data", label_key="lan_labels"): + """A minimal training-data pickle shaped like ssm-simulators output.""" + with open(path, "wb") as f: + pickle.dump( + { + features_key: np.random.randn(64, INPUT_DIM).astype(np.float32), + label_key: np.random.randn(64).astype(np.float32), + "generator_config": {"model": "ddm"}, + "model_config": {"params": ["v", "a", "z", "t"]}, + }, + f, + ) + return path + + +def _train_tiny_jax_network(tmp_path: Path) -> dict: + """Train a micro JaxMLP and return paths to its artifacts.""" + data_file = tmp_path / "training_data.pickle" + with open(data_file, "wb") as f: + pickle.dump( + { + "lan_data": np.random.randn(64, INPUT_DIM).astype(np.float32), + "lan_labels": np.random.randn(64).astype(np.float32), + "generator_config": {"model": "ddm"}, + "model_config": {"params": ["v", "a", "z", "t"]}, + }, + f, + ) + + dataset = lanfactory.trainers.DatasetTorch( + file_ids=[data_file], + batch_size=16, + features_key="lan_data", + label_key="lan_labels", + ) + dl = DataLoader(dataset, batch_size=None) + + net = lanfactory.trainers.JaxMLPFactory( + network_config=dict(NETWORK_CONFIG), train=True + ) + trainer = lanfactory.trainers.ModelTrainerJaxMLP( + train_config={ + "n_epochs": 2, # scheduler requires decay_steps > warmup_steps + "loss": "huber", + "optimizer": "adam", + "learning_rate": 0.001, + "lr_scheduler": None, + "lr_scheduler_params": {}, + "weight_decay": 0.0, + "train_output_type": "logprob", + }, + train_dl=dl, + valid_dl=dl, + model=net, + seed=42, + ) + out_dir = tmp_path / "nets" + state = trainer.train_and_evaluate( + output_folder=out_dir, + output_file_id="ddm", + run_id="runx", + mlflow_on=False, + save_outputs=True, + verbose=0, + network_type="lan", + ) + + def find(suffix): + matches = [p for p in out_dir.iterdir() if p.name.endswith(suffix)] + assert len(matches) == 1, (suffix, sorted(p.name for p in out_dir.iterdir())) + return matches[0] + + return { + "state_file": find("_train_state.jax"), + "onnx_file": find("_model.onnx"), + "trainer": trainer, + "state": state, + "out_dir": out_dir, + "tmp_path": tmp_path, + } + + +@pytest.fixture(scope="module") +def trained(tmp_path_factory): + return _train_tiny_jax_network(tmp_path_factory.mktemp("jax_export")) + + +class TestContract: + """The exported graph must satisfy the ecosystem ONNX contract.""" + + def test_all_input_dims_concrete(self, trained): + model = onnx.load(str(trained["onnx_file"])) + for inp in model.graph.input: + dims = inp.type.tensor_type.shape.dim + for d in dims: + assert d.HasField("dim_value"), ( + f"symbolic dim {d.dim_param!r} in {inp.name} — " + "dynamic axes are forbidden (HSSM rejects them at load)" + ) + + def test_input_shape_is_one_by_input_dim(self, trained): + model = onnx.load(str(trained["onnx_file"])) + dims = [d.dim_value for d in model.graph.input[0].type.tensor_type.shape.dim] + assert dims == [1, INPUT_DIM] + + def test_op_profile_is_mlp_only(self, trained): + # Gemm + elementwise ops only: the profile HSSM's rank-1 + vmap + # consumption is known to handle. Calibrated against what jax2onnx + # 0.15 actually emits: relu lowers to Max (not Relu), and the eval + # logsigmoid head of logits networks adds Neg/Exp/Add/Log. + model = onnx.load(str(trained["onnx_file"])) + ops = {n.op_type for n in model.graph.node} + assert ops <= ALLOWED_OPS, ops + + +class TestParity: + def test_ort_matches_jax_forward(self, trained): + """ONNX output == live jax forward, 1000 draws, float32 tolerance.""" + from functools import partial + + from lanfactory.trainers.jax_mlp import JaxMLP + + sess = ort.InferenceSession(str(trained["onnx_file"])) + iname = sess.get_inputs()[0].name + live = trained["trainer"].model + eval_model = JaxMLP( + layer_sizes=live.layer_sizes, + activations=live.activations, + train_output_type=live.train_output_type, + train=False, + ) + fwd = partial(eval_model.apply, trained["state"].params) + + import jax.numpy as jnp + + rng = np.random.default_rng(0) + X = rng.standard_normal((1000, INPUT_DIM)).astype(np.float32) + jax_out = np.asarray(fwd(jnp.asarray(X))) + max_err = 0.0 + for i in range(X.shape[0]): + o = sess.run(None, {iname: X[i : i + 1]})[0] + max_err = max(max_err, float(np.max(np.abs(o - jax_out[i])))) + assert max_err < 1e-4, f"parity violated: max|ORT - jax| = {max_err}" + + def test_file_based_transform_matches_in_trainer_export(self, trained): + """transform_jax_to_onnx on saved artifacts == the in-trainer export.""" + network_config_file = trained["tmp_path"] / "network_config.pickle" + with open(network_config_file, "wb") as f: + pickle.dump(dict(NETWORK_CONFIG), f) + + onnx_roundtrip = trained["tmp_path"] / "roundtrip.onnx" + transform_jax_to_onnx( + network_config_file=str(network_config_file), + state_file=str(trained["state_file"]), + input_shape=INPUT_DIM, + output_onnx_file=str(onnx_roundtrip), + ) + + sess_a = ort.InferenceSession(str(trained["onnx_file"])) + sess_b = ort.InferenceSession(str(onnx_roundtrip)) + rng = np.random.default_rng(1) + for _ in range(50): + x = rng.standard_normal((1, INPUT_DIM)).astype(np.float32) + oa = sess_a.run(None, {sess_a.get_inputs()[0].name: x})[0] + ob = sess_b.run(None, {sess_b.get_inputs()[0].name: x})[0] + np.testing.assert_allclose(oa, ob, atol=1e-6) + + +class TestTrainerFlag: + def test_no_export_onnx_skips_the_artifact(self, tmp_path): + artifacts = _train_tiny_jax_network(tmp_path) + # retrain into a fresh dir with export disabled + out_dir2 = tmp_path / "nets2" + artifacts["trainer"].train_and_evaluate( + output_folder=out_dir2, + output_file_id="ddm", + run_id="runy", + mlflow_on=False, + save_outputs=True, + verbose=0, + network_type="lan", + export_onnx=False, + ) + names = [p.name for p in out_dir2.iterdir()] + assert names, "no artifacts produced" + assert not any(n.endswith(".onnx") for n in names), names + + +class TestLogitsHead: + """CPN/OPN exports must emit the eval logsigmoid head, not raw logits. + + The raw-head export was a blocker: HSSM treats network output as + element-wise log-likelihood, so raw logits silently corrupt every logp + by +log(1+exp(-logit)). + """ + + def test_logits_export_equals_logsigmoid_of_raw_head(self, tmp_path): + from functools import partial + + import jax.numpy as jnp + from lanfactory.trainers.jax_mlp import JaxMLP + + f = make_training_pickle( + tmp_path / "training_data_cpn.pickle", + features_key="cpn_data", + label_key="cpn_labels", + ) + dataset = lanfactory.trainers.DatasetTorch( + file_ids=[f], batch_size=16, features_key="cpn_data", label_key="cpn_labels" + ) + dl = DataLoader(dataset, batch_size=None) + + config = { + "layer_sizes": [8, 1], + "activations": ["tanh", "linear"], + "train_output_type": "logits", + "network_type": "cpn", + } + net = lanfactory.trainers.JaxMLPFactory(network_config=config, train=True) + trainer = lanfactory.trainers.ModelTrainerJaxMLP( + train_config={ + "n_epochs": 2, + "loss": "bcelogit", + "optimizer": "adam", + "learning_rate": 0.001, + "lr_scheduler": None, + "lr_scheduler_params": {}, + "weight_decay": 0.0, + "train_output_type": "logits", + }, + train_dl=dl, + valid_dl=dl, + model=net, + seed=42, + ) + out_dir = tmp_path / "nets" + state = trainer.train_and_evaluate( + output_folder=out_dir, + output_file_id="ddm", + run_id="runc", + mlflow_on=False, + save_outputs=True, + verbose=0, + network_type="cpn", + ) + + onnx_files = [p for p in out_dir.iterdir() if p.suffix == ".onnx"] + assert len(onnx_files) == 1 + sess = ort.InferenceSession(str(onnx_files[0])) + iname = sess.get_inputs()[0].name + + raw_head = partial(trainer.model.apply, state.params) # train=True + eval_head = partial( + JaxMLP( + layer_sizes=trainer.model.layer_sizes, + activations=trainer.model.activations, + train_output_type="logits", + train=False, + ).apply, + state.params, + ) + + rng = np.random.default_rng(2) + for _ in range(100): + x = rng.standard_normal((1, INPUT_DIM)).astype(np.float32) + onnx_out = sess.run(None, {iname: x})[0] + raw = np.asarray(raw_head(jnp.asarray(x))) + expected = np.asarray(eval_head(jnp.asarray(x))) + # eval head == logsigmoid of raw head, and the export matches it + np.testing.assert_allclose( + expected, -np.log1p(np.exp(-raw)), rtol=1e-5, atol=1e-6 + ) + np.testing.assert_allclose(onnx_out, expected, rtol=1e-4, atol=1e-5) + # and decisively: the export is NOT the raw head + assert np.max(np.abs(onnx_out - raw)) > 1e-3