Correct metadata and docs our own merges made stale

CLAUDE.md

@@ -79,32 +79,46 @@ but differ in output type and loss function.
 - **PyTorch** (`torchtrain` CLI, `trainers/torch_mlp.py`) — primary backend.
   Supports CUDA, ONNX export, full training loop with validation.
 - **JAX/Flax** (`jaxtrain` CLI, `trainers/jax_mlp.py`) — alternative backend.
-  Uses optax optimizers. No native ONNX export (train in JAX, convert via PyTorch if needed).
+  Uses optax optimizers. Exports ONNX directly via `jax2onnx` (`onnx/jax_export.py`);
+  `jaxtrain` writes the artifact by default (`--no-export-onnx` to skip).
 
 ### ONNX Export Pipeline
 
-Three exporters (all in `src/lanfactory/onnx/`) produce `.onnx` via `torch.onnx.export()`:
-- **LAN/CPN/OPN MLPs** — `transform_onnx.py` (`transform-onnx` CLI)
+Four exporters, all in `src/lanfactory/onnx/`:
+- **LAN/CPN/OPN MLPs (torch)** — `transform_onnx.py` (`transform-onnx` CLI)
+- **LAN/CPN/OPN MLPs (jax)** — `jax_export.py` (`transform-jax-onnx` CLI), via `jax2onnx`
 - **sbi** posterior/likelihood/ratio estimators — `sbi.py` (`transform_sbi_to_onnx`)
 - **bayesflow** networks — `bayesflow.py` (`transform_bayesflow_to_onnx`)
 
-All follow the single-trial contract: export with a concrete rank-1 per-trial
-input shape (no `dynamic_axes`); HSSM batches per-trial via `jax.vmap`. This is
-the format HSSM consumes at runtime.
+All follow the single-trial contract: **every input dim concrete, no
+`dynamic_axes`**; HSSM batches per-trial via `jax.vmap`.
+
+The *rank* is not part of the contract — it follows from how your tracer lowers
+a dense layer. The MLP exporters (torch and jax) trace `(1, D)` and lower to
+`Gemm`, whose ONNX spec requires rank 2; the sbi and bayesflow exporters trace
+rank-1 `(D,)` because `torch.onnx.export` lowers `Linear` to rank-agnostic
+`MatMul`+`Add`. Both load in HSSM and, measured under `vmap`+`jit`, run
+identically. The production networks on franklab/HSSM are `(1, D)` Gemm.
+`assert_single_trial_contract` in `onnx/contract.py` is the executable version
+of this paragraph — call it from any new exporter's tests.
 
 ### HuggingFace Integration
 
 - **Upload:** `lanfactory.hf.upload_model()` — uploads `.onnx`, `.pt`, config pickles,
-  and auto-generated README to `franklab/HSSM` on HuggingFace.
-  Requires `model_card.yaml` in the model folder.
+  and auto-generated README to `franklab/HSSM` on HuggingFace. Publishes the
+  canonical ONNX at the repo *root* under the filename HSSM downloads, plus the
+  full artifact set under `{network_type}/{model}/`, plus a root `manifest.json`
+  — in one atomic commit. Refuses to replace an existing root network without
+  `--overwrite-root`. `model_card.yaml` is generated when absent
+  (`--require-model-card` to demand one).
 - **Download:** `lanfactory.hf.download_model()` — downloads by network type + model name.
 - **Default repo:** `franklab/HSSM`
 - **Optional dependency:** `huggingface-hub>=0.20.0` (install via `uv sync --extra hf`)
 
 ### Config System
 
 Training configs are YAML files parsed by the CLI. Key fields:
-- `NETWORK_TYPE`: `lan`, `cpn`, or `opn`
+- `NETWORK_TYPE`: `lan`, `cpn`, `opn`, or `gonogo`
 - `layer_sizes`, `activations`: network architecture
 - `n_epochs`, `learning_rate`, `loss`, `optimizer`: training hyperparams
 - `cpu_batch_size`, `gpu_batch_size`: device-specific batch sizes
@@ -123,6 +137,7 @@ Optional experiment tracking via MLflow. CLI flags: `--mlflow-run-name`, `--mlfl
 | `torchtrain` | `lanfactory.cli.torch_train` | Train PyTorch networks from YAML config |
 | `jaxtrain` | `lanfactory.cli.jax_train` | Train JAX networks from YAML config |
 | `transform-onnx` | `lanfactory.onnx.transform_onnx` | Convert PyTorch model → ONNX |
+| `transform-jax-onnx` | `lanfactory.onnx.jax_export` | Convert a jaxtrain network → ONNX |
 | `upload-hf` | `lanfactory.cli.upload_hf` | Upload trained models to HuggingFace |
 | `download-hf` | `lanfactory.cli.download_hf` | Download models from HuggingFace |
 

src/lanfactory/cli/download_hf.py

@@ -23,7 +23,7 @@ def main(
     network_type: str = typer.Option(
         ...,
         "--network-type",
-        help="Network type: lan, cpn, or opn.",
+        help=f"Network type: one of {', '.join(VALID_NETWORK_TYPES)}.",
     ),
     model_name: str = typer.Option(
         ...,
@@ -116,7 +116,8 @@ def main(
     except ImportError as e:
         logger.error(
             "huggingface_hub is required for HuggingFace downloads. "
-            "Install it with: pip install lanfactory[hf]"
+            "Install it with: pip install 'lanfactory[hf]' "
+            "(or, in a checkout: uv sync --extra hf)."
         )
         raise typer.Exit(code=1) from e
 

src/lanfactory/cli/upload_hf.py

@@ -32,7 +32,7 @@ def main(
     network_type: str = typer.Option(
         ...,
         "--network-type",
-        help="Network type: lan, cpn, or opn.",
+        help=f"Network type: one of {', '.join(VALID_NETWORK_TYPES)}.",
     ),
     model_name: str = typer.Option(
         ...,
@@ -174,7 +174,8 @@ def main(
     except ImportError as e:
         logger.error(
             "huggingface_hub is required for HuggingFace uploads. "
-            "Install it with: pip install lanfactory[hf]"
+            "Install it with: pip install 'lanfactory[hf]' "
+            "(or, in a checkout: uv sync --extra hf)."
         )
         raise typer.Exit(code=1) from e
 

src/lanfactory/hf/__init__.py

@@ -5,6 +5,10 @@
 """
 
 DEFAULT_REPO_ID = "franklab/HSSM"
+# Matches what franklab/HSSM actually declares. The *code* in this ecosystem is
+# MIT, but a model card describes the published artifact, so it follows the
+# artifact repo — auto-generated cards claiming MIT would contradict it.
+DEFAULT_LICENSE = "bsd-2-clause"
 # gonogo included: the trainers already build gonogo networks (cli/utils.py
 # train_output_type_dict) and HSSM resolves "{model}_gonogo.onnx", so excluding
 # it here made a trainable, loadable network type unpublishable.
@@ -20,6 +24,7 @@
 
 __all__ = [
     "DEFAULT_REPO_ID",
+    "DEFAULT_LICENSE",
     "VALID_NETWORK_TYPES",
     "load_model_card_yaml",
     "generate_readme",

src/lanfactory/hf/download.py

@@ -103,7 +103,8 @@ def _download_model_hf(  # pragma: no cover
     except ImportError as exc:
         raise ImportError(
             "huggingface_hub is required for HuggingFace downloads. "
-            "Install it with: pip install lanfactory[hf]"
+            "Install it with: pip install 'lanfactory[hf]' "
+            "(or, in a checkout: uv sync --extra hf)."
         ) from exc
 
     output_folder.mkdir(parents=True, exist_ok=True)

src/lanfactory/hf/model_card.py

@@ -12,6 +12,8 @@
 
 import yaml
 
+from lanfactory.hf import DEFAULT_LICENSE
+
 logger = logging.getLogger(__name__)
 
 
@@ -26,7 +28,8 @@ class ModelCardConfig:
     library_name : str
         Library name for HuggingFace (default: "onnx").
     license : str
-        License identifier (default: "mit").
+        License identifier (defaults to DEFAULT_LICENSE, which tracks the
+        artifact repo franklab/HSSM — currently bsd-2-clause).
     title : str
         Model title.
     description : str
@@ -41,7 +44,7 @@ class ModelCardConfig:
 
     tags: list[str] = field(default_factory=lambda: ["lan", "ssm", "hssm"])
     library_name: str = "onnx"
-    license: str = "mit"
+    license: str = DEFAULT_LICENSE
     title: str = "LAN Model"
     description: str = "Likelihood Approximation Network trained with LANfactory."
     architecture: dict | None = None
@@ -82,7 +85,7 @@ def load_model_card_yaml(model_folder: Path) -> ModelCardConfig:
     config = ModelCardConfig(
         tags=data.get("tags", ["lan", "ssm", "hssm"]),
         library_name=data.get("library_name", "onnx"),
-        license=data.get("license", "mit"),
+        license=data.get("license", DEFAULT_LICENSE),
         title=data.get("title", "LAN Model"),
         description=data.get(
             "description", "Likelihood Approximation Network trained with LANfactory."

src/lanfactory/hf/upload.py

@@ -10,7 +10,7 @@
 import tempfile
 from pathlib import Path
 
-from lanfactory.hf import DEFAULT_REPO_ID, VALID_NETWORK_TYPES
+from lanfactory.hf import DEFAULT_LICENSE, DEFAULT_REPO_ID, VALID_NETWORK_TYPES
 
 logger = logging.getLogger(__name__)
 
@@ -424,7 +424,7 @@ def write_default_model_card(
     card = {
         "tags": [network_type, "ssm", "hssm"],
         "library_name": "onnx",
-        "license": "mit",
+        "license": DEFAULT_LICENSE,
         "title": f"{model_name} ({network_type.upper()})",
         "description": (
             f"{network_type.upper()} for the {model_name} sequential sampling "
@@ -531,7 +531,8 @@ def _upload_to_hf(  # pragma: no cover
     except ImportError as exc:
         raise ImportError(
             "huggingface_hub is required for HuggingFace uploads. "
-            "Install it with: pip install lanfactory[hf]"
+            "Install it with: pip install 'lanfactory[hf]' "
+            "(or, in a checkout: uv sync --extra hf)."
         ) from exc
 
     api = HfApi(token=token)

src/lanfactory/onnx/__init__.py

@@ -3,9 +3,12 @@
 from .sbi import transform_sbi_to_onnx
 from .transform_onnx import transform_to_onnx
 
+from lanfactory.onnx.contract import assert_single_trial_contract  # noqa: E402
+
 __all__ = [
     "transform_to_onnx",
     "transform_jax_to_onnx",
     "transform_sbi_to_onnx",
     "transform_bayesflow_to_onnx",
+    "assert_single_trial_contract",
 ]

src/lanfactory/onnx/contract.py

@@ -0,0 +1,88 @@
+"""The single-trial ONNX contract, as a check instead of a paragraph.
+
+Every exporter here produces artifacts that HSSM loads through
+``jaxonnxruntime``, which traces against the construction-time dummy and bakes
+the resulting shapes into the returned closure. A graph with a dynamic axis
+therefore does not fail loudly — it silently returns wrong numbers for any
+model with a batch-dependent intermediate. HSSM guards its own door
+(``make_jax_func`` raises on symbolic dims), but by then the artifact is
+published.
+
+The invariant is exactly one thing: **every input dimension is concrete**.
+
+Rank is *not* part of it, which is the part that keeps getting misremembered.
+It follows from how a tracer lowers a dense layer: ``torch.onnx.export`` on a
+rank-1 dummy gives rank-agnostic ``MatMul``+``Add`` (sbi, bayesflow), while a
+``(1, D)`` dummy — and ``jax2onnx`` always — gives ``Gemm``, whose ONNX spec
+*requires* rank 2. Both load in HSSM and run identically once XLA has fused
+them. The production networks on franklab/HSSM are ``(1, D)`` Gemm.
+
+Call this from an exporter's tests rather than restating the rules.
+"""
+
+from pathlib import Path
+
+
+def assert_single_trial_contract(
+    onnx_path: str | Path,
+    expected_input_width: int | None = None,
+    allowed_ops: set[str] | None = None,
+) -> dict:
+    """Raise AssertionError unless the artifact satisfies the contract.
+
+    Parameters
+    ----------
+    onnx_path
+        The exported artifact.
+    expected_input_width
+        The per-trial input width (the last dimension), when the caller knows
+        it. Catches an exporter that silently changed its input layout.
+    allowed_ops
+        When given, the graph's op types must be a subset. Useful to pin a
+        lowering that a pinned 0.x exporter dependency could change under you.
+
+    Returns
+    -------
+    dict
+        ``{"input_shape", "input_width", "ops"}`` for further assertions.
+    """
+    import onnx
+    import onnxruntime as ort
+
+    onnx_path = Path(onnx_path)
+    model = onnx.load(str(onnx_path))
+    onnx.checker.check_model(model)
+
+    for graph_input in model.graph.input:
+        for dim in graph_input.type.tensor_type.shape.dim:
+            assert dim.HasField("dim_value"), (
+                f"symbolic dim {dim.dim_param!r} in input {graph_input.name!r}: "
+                "HSSM's make_jax_func rejects dynamic axes at load, and a graph "
+                "that slips through returns wrong numbers rather than failing"
+            )
+            # A dim_value of 0 is set-but-not-concrete: some producers use it
+            # for "unknown", and it is a degenerate axis either way.
+            assert dim.dim_value > 0, (
+                f"zero dim in input {graph_input.name!r}: not a concrete shape"
+            )
+
+    # onnxruntime is the arbiter of whether the graph is actually runnable:
+    # a rank-1-traced Gemm passes the checker and fails here.
+    session = ort.InferenceSession(str(onnx_path))
+    input_shape = session.get_inputs()[0].shape
+    input_width = int(input_shape[-1])
+
+    if expected_input_width is not None:
+        assert input_width == expected_input_width, (
+            f"input width {input_width} != expected {expected_input_width}"
+        )
+
+    ops = {node.op_type for node in model.graph.node}
+    if allowed_ops is not None:
+        assert ops <= allowed_ops, f"unexpected ops {sorted(ops - allowed_ops)}"
+
+    return {
+        "input_shape": list(input_shape),
+        "input_width": input_width,
+        "ops": sorted(ops),
+    }

tests/hf/test_dual_layout.py

@@ -748,3 +748,18 @@ def spy_repo_info(self, repo_id, revision=None):
         # and the manifest is read AT that parent, not at the moving branch
         assert dict(calls)["fetch_manifest"] == "sha1"
         assert commits[0]["parent_commit"] == "sha1"
+
+
+def test_generated_card_uses_the_artifact_repo_license(tmp_path):
+    """franklab/HSSM declares bsd-2-clause; the code is MIT but a model card
+    describes the artifact, so a generated card must not contradict the repo."""
+    import yaml
+
+    from lanfactory.hf import DEFAULT_LICENSE
+    from lanfactory.hf.upload import write_default_model_card
+
+    # One literal pin of the value, then wiring checked against the constant.
+    assert DEFAULT_LICENSE == "bsd-2-clause"
+    write_default_model_card(tmp_path, "lan", "ddm")
+    card = yaml.safe_load((tmp_path / "model_card.yaml").read_text())
+    assert card["license"] == DEFAULT_LICENSE

tests/hf/test_model_card.py

@@ -5,6 +5,7 @@
 import pytest
 import yaml
 
+from lanfactory.hf import DEFAULT_LICENSE
 from lanfactory.hf.model_card import (
     ModelCardConfig,
     generate_readme,
@@ -21,7 +22,9 @@ def test_default_values(self):
         config = ModelCardConfig()
         assert config.tags == ["lan", "ssm", "hssm"]
         assert config.library_name == "onnx"
-        assert config.license == "mit"
+        # Follows the artifact repo (franklab/HSSM is bsd-2-clause), not the
+        # ecosystem's code licence — a card describes the published artifact.
+        assert config.license == "bsd-2-clause"
         assert config.title == "LAN Model"
         assert config.architecture is None
         assert config.training is None
@@ -131,7 +134,9 @@ def test_generates_valid_frontmatter(self):
 
         assert frontmatter["tags"] == ["lan", "ssm", "ddm"]
         assert frontmatter["library_name"] == "onnx"
-        assert frontmatter["license"] == "mit"
+        # The config did not set a licence, so the frontmatter carries the
+        # default — which tracks the artifact repo, not the code licence.
+        assert frontmatter["license"] == DEFAULT_LICENSE
 
     def test_includes_title_and_description(self):
         """Test that README includes title and description."""