From bbd8ae55769df793a315cd3d125c3962aa14ca73 Mon Sep 17 00:00:00 2001 From: Alexander Fengler Date: Sun, 9 Aug 2026 22:04:08 -0400 Subject: [PATCH 1/2] chore: correct metadata and docs our own merges made stale MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Model cards defaulted to license 'mit', but franklab/HSSM declares bsd-2-clause. The code is MIT; a card describes the *artifact*, so it follows the artifact repo. One DEFAULT_LICENSE next to DEFAULT_REPO_ID. Verified against every published artifact before the first real publish. - The 'install it with' hints named only pip. Reworded to give both, since the message reaches end users (who pip-installed) and developers (who did not). The repo's own uv rule carves out end-user install docs. - CLI --network-type help still said 'lan, cpn, or opn' after gonogo was added; it now derives from VALID_NETWORK_TYPES so it cannot drift again. - CLAUDE.md claimed jaxtrain has no ONNX export (false since the jax2onnx exporter), that model_card.yaml is required (false since it is generated), and that all exporters trace rank-1 (never true of the MLP exporter, and not true of the production artifacts). Adds onnx/contract.py: assert_single_trial_contract, the executable form of the rank paragraph. The invariant is concrete input dims; rank follows from the tracer's lowering. Checked against all 18 production artifacts on franklab/HSSM — every one passes, all (1, D) Gemm+Tanh — so this encodes what production already is rather than imposing something new. Co-Authored-By: Claude Opus 5 --- CLAUDE.md | 33 ++++++++---- src/lanfactory/cli/download_hf.py | 5 +- src/lanfactory/cli/upload_hf.py | 5 +- src/lanfactory/hf/__init__.py | 4 ++ src/lanfactory/hf/download.py | 3 +- src/lanfactory/hf/model_card.py | 4 +- src/lanfactory/hf/upload.py | 7 +-- src/lanfactory/onnx/__init__.py | 1 + src/lanfactory/onnx/contract.py | 83 +++++++++++++++++++++++++++++++ tests/hf/test_dual_layout.py | 14 ++++++ tests/onnx/test_contract.py | 60 ++++++++++++++++++++++ 11 files changed, 201 insertions(+), 18 deletions(-) create mode 100644 src/lanfactory/onnx/contract.py create mode 100644 tests/onnx/test_contract.py diff --git a/CLAUDE.md b/CLAUDE.md index 1961ed9..65340f9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -79,24 +79,38 @@ 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`) @@ -104,7 +118,7 @@ the format HSSM consumes at runtime. ### 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 | diff --git a/src/lanfactory/cli/download_hf.py b/src/lanfactory/cli/download_hf.py index 262f776..8ab776d 100644 --- a/src/lanfactory/cli/download_hf.py +++ b/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 diff --git a/src/lanfactory/cli/upload_hf.py b/src/lanfactory/cli/upload_hf.py index 057cb9c..849233a 100644 --- a/src/lanfactory/cli/upload_hf.py +++ b/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 diff --git a/src/lanfactory/hf/__init__.py b/src/lanfactory/hf/__init__.py index 4feb3fd..54bf482 100644 --- a/src/lanfactory/hf/__init__.py +++ b/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. diff --git a/src/lanfactory/hf/download.py b/src/lanfactory/hf/download.py index 52bb17c..3b3a837 100644 --- a/src/lanfactory/hf/download.py +++ b/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) diff --git a/src/lanfactory/hf/model_card.py b/src/lanfactory/hf/model_card.py index 55bdf3e..b27363d 100644 --- a/src/lanfactory/hf/model_card.py +++ b/src/lanfactory/hf/model_card.py @@ -12,6 +12,8 @@ import yaml +from lanfactory.hf import DEFAULT_LICENSE + logger = logging.getLogger(__name__) @@ -82,7 +84,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." diff --git a/src/lanfactory/hf/upload.py b/src/lanfactory/hf/upload.py index fd19870..3d0f3d2 100644 --- a/src/lanfactory/hf/upload.py +++ b/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) diff --git a/src/lanfactory/onnx/__init__.py b/src/lanfactory/onnx/__init__.py index 626293d..f3fbcc2 100755 --- a/src/lanfactory/onnx/__init__.py +++ b/src/lanfactory/onnx/__init__.py @@ -9,3 +9,4 @@ "transform_sbi_to_onnx", "transform_bayesflow_to_onnx", ] +from lanfactory.onnx.contract import assert_single_trial_contract # noqa: E402,F401 diff --git a/src/lanfactory/onnx/contract.py b/src/lanfactory/onnx/contract.py new file mode 100644 index 0000000..5ee1b15 --- /dev/null +++ b/src/lanfactory/onnx/contract.py @@ -0,0 +1,83 @@ +"""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" + ) + + # 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), + } diff --git a/tests/hf/test_dual_layout.py b/tests/hf/test_dual_layout.py index 920f9e2..1669b1f 100644 --- a/tests/hf/test_dual_layout.py +++ b/tests/hf/test_dual_layout.py @@ -748,3 +748,17 @@ 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 + + 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"] == "bsd-2-clause" diff --git a/tests/onnx/test_contract.py b/tests/onnx/test_contract.py new file mode 100644 index 0000000..ba3ce94 --- /dev/null +++ b/tests/onnx/test_contract.py @@ -0,0 +1,60 @@ +"""Tests for the executable single-trial ONNX contract.""" + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper + +from lanfactory.onnx import assert_single_trial_contract + + +def make_onnx(path, input_dims): + width = input_dims[-1] if isinstance(input_dims[-1], int) else 6 + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, list(input_dims)) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1]) + w = helper.make_tensor( + "w", TensorProto.FLOAT, [width, 1], np.zeros(width, dtype=np.float32).tolist() + ) + graph = helper.make_graph( + [helper.make_node("MatMul", ["x", "w"], ["y"])], "g", [x], [y], initializer=[w] + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 14)]) + model.ir_version = 8 + onnx.save(model, str(path)) + return path + + +def test_accepts_a_concrete_graph(tmp_path): + result = assert_single_trial_contract( + make_onnx(tmp_path / "ok.onnx", (1, 6)), expected_input_width=6 + ) + assert result["input_shape"] == [1, 6] + + +def test_rejects_a_symbolic_dim(tmp_path): + # The whole point: jaxonnxruntime bakes shapes at trace time, so a dynamic + # axis returns wrong numbers rather than failing. + with pytest.raises(AssertionError, match="symbolic dim"): + assert_single_trial_contract(make_onnx(tmp_path / "dyn.onnx", ("batch", 6))) + + +def test_rejects_an_unexpected_input_width(tmp_path): + with pytest.raises(AssertionError, match="input width"): + assert_single_trial_contract( + make_onnx(tmp_path / "w.onnx", (1, 5)), expected_input_width=6 + ) + + +def test_rejects_ops_outside_the_allowed_set(tmp_path): + # Guards against a pinned 0.x exporter changing its lowering under us. + with pytest.raises(AssertionError, match="unexpected ops"): + assert_single_trial_contract( + make_onnx(tmp_path / "ops.onnx", (1, 6)), allowed_ops={"Gemm", "Tanh"} + ) + + +def test_rank_is_not_part_of_the_contract(tmp_path): + """sbi and bayesflow trace rank-1; the MLP exporters trace (1, D). Both are + valid — only concreteness is required.""" + assert assert_single_trial_contract(make_onnx(tmp_path / "r1.onnx", (6,))) + assert assert_single_trial_contract(make_onnx(tmp_path / "r2.onnx", (1, 6))) From 42602a669d46e4d08114ba93b224e3db5ad35b24 Mon Sep 17 00:00:00 2001 From: Alexander Fengler Date: Mon, 10 Aug 2026 20:19:01 -0400 Subject: [PATCH 2/2] review: close the license default at its source; tighten the contract CodeRabbit was right that my license fix was incomplete. I changed the YAML fallback and the generated card, but ModelCardConfig's own field still defaulted to 'mit', so any direct caller of ModelCardConfig() would still have minted a card contradicting the artifact repo. It now defaults to DEFAULT_LICENSE, and two pre-existing tests that pinned 'mit' are updated to the corrected contract (one keeps an explicit 'mit' to prove an explicit licence still passes through untouched). Copilot findings, all valid: - assert_single_trial_contract accepted a dim_value of 0. HasField is true for an explicitly-set zero, which some producers use for 'unknown' and which is a degenerate axis regardless. Now rejected, with a test. - DEFAULT_LICENSE and assert_single_trial_contract were importable but missing from __all__. - The contract test helper declared a [1, 1] output for the rank-1 input case, where MatMul actually yields rank-1. Leaning on a permissive checker in the tests for the checker is not a good look. - The license test hard-coded the string twice; it now pins the literal once and checks the wiring against the constant. Co-Authored-By: Claude Opus 5 --- src/lanfactory/hf/__init__.py | 1 + src/lanfactory/hf/model_card.py | 5 +++-- src/lanfactory/onnx/__init__.py | 4 +++- src/lanfactory/onnx/contract.py | 5 +++++ tests/hf/test_dual_layout.py | 3 ++- tests/hf/test_model_card.py | 9 +++++++-- tests/onnx/test_contract.py | 20 +++++++++++++++++++- 7 files changed, 40 insertions(+), 7 deletions(-) diff --git a/src/lanfactory/hf/__init__.py b/src/lanfactory/hf/__init__.py index 54bf482..95e94b0 100644 --- a/src/lanfactory/hf/__init__.py +++ b/src/lanfactory/hf/__init__.py @@ -24,6 +24,7 @@ __all__ = [ "DEFAULT_REPO_ID", + "DEFAULT_LICENSE", "VALID_NETWORK_TYPES", "load_model_card_yaml", "generate_readme", diff --git a/src/lanfactory/hf/model_card.py b/src/lanfactory/hf/model_card.py index b27363d..ac5ecf8 100644 --- a/src/lanfactory/hf/model_card.py +++ b/src/lanfactory/hf/model_card.py @@ -28,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 @@ -43,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 diff --git a/src/lanfactory/onnx/__init__.py b/src/lanfactory/onnx/__init__.py index f3fbcc2..f3452f4 100755 --- a/src/lanfactory/onnx/__init__.py +++ b/src/lanfactory/onnx/__init__.py @@ -3,10 +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", ] -from lanfactory.onnx.contract import assert_single_trial_contract # noqa: E402,F401 diff --git a/src/lanfactory/onnx/contract.py b/src/lanfactory/onnx/contract.py index 5ee1b15..217663e 100644 --- a/src/lanfactory/onnx/contract.py +++ b/src/lanfactory/onnx/contract.py @@ -60,6 +60,11 @@ def assert_single_trial_contract( "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. diff --git a/tests/hf/test_dual_layout.py b/tests/hf/test_dual_layout.py index 1669b1f..172868a 100644 --- a/tests/hf/test_dual_layout.py +++ b/tests/hf/test_dual_layout.py @@ -758,7 +758,8 @@ def test_generated_card_uses_the_artifact_repo_license(tmp_path): 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"] == "bsd-2-clause" + assert card["license"] == DEFAULT_LICENSE diff --git a/tests/hf/test_model_card.py b/tests/hf/test_model_card.py index 82e3525..075cfd6 100644 --- a/tests/hf/test_model_card.py +++ b/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.""" diff --git a/tests/onnx/test_contract.py b/tests/onnx/test_contract.py index ba3ce94..75ba763 100644 --- a/tests/onnx/test_contract.py +++ b/tests/onnx/test_contract.py @@ -10,8 +10,11 @@ def make_onnx(path, input_dims): width = input_dims[-1] if isinstance(input_dims[-1], int) else 6 + # MatMul drops the leading axis for a rank-1 input, so declare the shape + # the graph actually produces rather than leaning on a permissive checker. + output_dims = [1, 1] if len(input_dims) == 2 else [1] x = helper.make_tensor_value_info("x", TensorProto.FLOAT, list(input_dims)) - y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, output_dims) w = helper.make_tensor( "w", TensorProto.FLOAT, [width, 1], np.zeros(width, dtype=np.float32).tolist() ) @@ -58,3 +61,18 @@ def test_rank_is_not_part_of_the_contract(tmp_path): valid — only concreteness is required.""" assert assert_single_trial_contract(make_onnx(tmp_path / "r1.onnx", (6,))) assert assert_single_trial_contract(make_onnx(tmp_path / "r2.onnx", (1, 6))) + + +def test_rejects_a_zero_dim(tmp_path): + """A dim_value of 0 is set-but-not-concrete — some producers use it for + "unknown" — so HasField alone is not enough.""" + path = make_onnx(tmp_path / "zero.onnx", (0, 6)) + with pytest.raises(AssertionError, match="zero dim"): + assert_single_trial_contract(path) + + +def test_model_card_config_default_license_is_not_mit(): + """A direct ModelCardConfig() must not mint a card contradicting the repo.""" + from lanfactory.hf import DEFAULT_LICENSE, ModelCardConfig + + assert ModelCardConfig().license == DEFAULT_LICENSE