jax → ONNX export via jax2onnx: jaxtrain now produces HSSM-consumable artifacts

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"
 

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,
     )
     # -------------------------------------------------------------
 

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",
 ]

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()

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}")